[mypyc] feat: new primitive for int.to_bytes

Author: BobTheBuidler

mypyc/lib-rt/int_ops.c

@@ -581,3 +581,54 @@ double CPyTagged_TrueDivide(CPyTagged x, CPyTagged y) {
     }
     return 1.0;
 }
+
+// int.to_bytes(length, byteorder, signed=False)
+PyObject *CPyInt_ToBytes(CPyTagged self, Py_ssize_t length, PyObject *byteorder, int signed_flag) {
+    PyObject *pyint = CPyTagged_StealAsObject(self);
+    if (!PyLong_Check(pyint)) {
+        Py_DECREF(pyint);
+        PyErr_SetString(PyExc_TypeError, "self must be int");
+        return NULL;
+    }
+    if (!PyUnicode_Check(byteorder)) {
+        Py_DECREF(pyint);
+        PyErr_SetString(PyExc_TypeError, "byteorder must be str");
+        return NULL;
+    }
+    const char *order = PyUnicode_AsUTF8(byteorder);
+    if (!order) {
+        Py_DECREF(pyint);
+        return NULL;
+    }
+    PyObject *result = PyLong_ToBytes(pyint, length, order, signed_flag);
+    Py_DECREF(pyint);
+    return result;
+}
+
+// Helper for PyLong_ToBytes (Python 3.2+)
+PyObject *PyLong_ToBytes(PyObject *v, Py_ssize_t length, const char *byteorder, int signed_flag) {
+    // This is a wrapper for PyLong_AsByteArray and PyBytes_FromStringAndSize
+    unsigned char *bytes = (unsigned char *)PyMem_Malloc(length);
+    if (!bytes) {
+        PyErr_NoMemory();
+        return NULL;
+    }
+    int little_endian = 0;
+    if (strcmp(byteorder, "little") == 0) {
+        little_endian = 1;
+    } else if (strcmp(byteorder, "big") == 0) {
+        little_endian = 0;
+    } else {
+        PyMem_Free(bytes);
+        PyErr_SetString(PyExc_ValueError, "byteorder must be either 'little' or 'big'");
+        return NULL;
+    }
+    int res = PyLong_AsByteArray((PyLongObject *)v, bytes, length, little_endian, signed_flag);
+    if (res < 0) {
+        PyMem_Free(bytes);
+        return NULL;
+    }
+    PyObject *result = PyBytes_FromStringAndSize((const char *)bytes, length);
+    PyMem_Free(bytes);
+    return result;
+}

mypyc/primitives/int_ops.py

@@ -21,6 +21,7 @@
     RType,
     bit_rprimitive,
     bool_rprimitive,
+    bytes_rprimitive,
     c_pyssize_t_rprimitive,
     float_rprimitive,
     int16_rprimitive,
@@ -31,7 +32,7 @@
     str_rprimitive,
     void_rtype,
 )
-from mypyc.primitives.registry import binary_op, custom_op, function_op, load_address_op, unary_op
+from mypyc.primitives.registry import binary_op, custom_op, function_op, load_address_op, method_op, unary_op
 
 # Constructors for builtins.int and native int types have the same behavior. In
 # interpreted mode, native int types are just aliases to 'int'.
@@ -305,3 +306,21 @@ def int_unary_op(name: str, c_function_name: str) -> PrimitiveDescription:
     c_function_name="PyLong_Check",
     error_kind=ERR_NEVER,
 )
+
+# int.to_bytes(length, byteorder, *, signed=False)
+method_op(
+    name="to_bytes",
+    arg_types=[int_rprimitive, int_rprimitive, str_rprimitive, bool_rprimitive],
+    return_type=bytes_rprimitive,
+    c_function_name="CPyInt_ToBytes",
+    error_kind=ERR_MAGIC,
+)
+
+# int.bit_length()
+method_op(
+    name="bit_length",
+    arg_types=[int_rprimitive],
+    return_type=int_rprimitive,
+    c_function_name="CPyInt_BitLength",
+    error_kind=ERR_MAGIC,
+)

mypyc/test-data/irbuild-int.test

@@ -210,3 +210,14 @@ L0:
     r0 = CPyTagged_Invert(n)
     x = r0
     return x
+
+[case testIntToBytes]
+def f(x: int) -> bytes:
+    return x.to_bytes(2, "big")
+[out]
+def f(x):
+    x :: int
+    r0 :: bytes
+L0:
+    r0 = int_to_bytes x, 2, 'big', 0
+    return r0

mypyc/test-data/run-integers.test

@@ -572,3 +572,12 @@ class subc(int):
 [file userdefinedint.py]
 class int:
     pass
+
+[case testIntToBytes]
+def to_bytes(n: int, length: int, byteorder: str, signed: bool = False) -> bytes:
+    return n.to_bytes(length, byteorder, signed=signed)
+def test_to_bytes() -> None:
+    assert to_bytes(255, 2, "big") == b'\x00\xff'
+    assert to_bytes(255, 2, "little") == b'\xff\x00'
+    assert to_bytes(-1, 2, "big", True) == b'\xff\xff'
+    assert to_bytes(0, 1, "big") == b'\x00'