[mypyc] feat: new primitive for int.to_bytes

mypyc/lib-rt/CPy.h

@@ -139,6 +139,8 @@ CPyTagged CPyTagged_Remainder_(CPyTagged left, CPyTagged right);
 CPyTagged CPyTagged_BitwiseLongOp_(CPyTagged a, CPyTagged b, char op);
 CPyTagged CPyTagged_Rshift_(CPyTagged left, CPyTagged right);
 CPyTagged CPyTagged_Lshift_(CPyTagged left, CPyTagged right);
+PyObject *CPyTagged_ToBytes(CPyTagged self, Py_ssize_t length, PyObject *byteorder, int signed_flag);
+PyObject *CPyLong_ToBytes(PyObject *v, Py_ssize_t length, const char *byteorder, int signed_flag);
 
 PyObject *CPyTagged_Str(CPyTagged n);
 CPyTagged CPyTagged_FromFloat(double f);

mypyc/lib-rt/int_ops.c

@@ -581,3 +581,59 @@ double CPyTagged_TrueDivide(CPyTagged x, CPyTagged y) {
     }
     return 1.0;
 }
+
+// int.to_bytes(length, byteorder, signed=False)
+PyObject *CPyTagged_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 = CPyLong_ToBytes(pyint, length, order, signed_flag);
+    Py_DECREF(pyint);
+    return result;
+}
+
+
+// Helper for CPyLong_ToBytes (Python 3.2+)
+PyObject *CPyLong_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;
+    if (strcmp(byteorder, "big") == 0) {
+        little_endian = 0;
+    } else if (strcmp(byteorder, "little") == 0) {
+        little_endian = 1;
+    } else {
+        PyMem_Free(bytes);
+        PyErr_SetString(PyExc_ValueError, "byteorder must be either 'little' or 'big'");
+        return NULL;
+    }
+#if PY_VERSION_HEX >= 0x030D0000  // 3.13.0
+    int res = _PyLong_AsByteArray((PyLongObject *)v, bytes, length, little_endian, signed_flag, 1);
+#else
+    int res = _PyLong_AsByteArray((PyLongObject *)v, bytes, length, little_endian, signed_flag);
+#endif
+    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,14 @@
     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 +313,22 @@ 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)
+method_op(
+    name="to_bytes",
+    arg_types=[int_rprimitive, int_rprimitive, str_rprimitive],
+    extra_int_constants=[(0, bool_rprimitive)],
+    return_type=bytes_rprimitive,
+    c_function_name="CPyTagged_ToBytes",
+    error_kind=ERR_MAGIC,
+)
+
+# int.to_bytes(length, byteorder, signed)
+method_op(
+    name="to_bytes",
+    arg_types=[int_rprimitive, int_rprimitive, str_rprimitive, bool_rprimitive],
+    return_type=bytes_rprimitive,
+    c_function_name="CPyTagged_ToBytes",
+    error_kind=ERR_MAGIC,
+)

mypyc/test-data/fixtures/ir.py

@@ -4,7 +4,7 @@
 import _typeshed
 from typing import (
     TypeVar, Generic, List, Iterator, Iterable, Dict, Optional, Tuple, Any, Set,
-    overload, Mapping, Union, Callable, Sequence, FrozenSet, Protocol
+    overload, Mapping, Union, Callable, Sequence, FrozenSet, Protocol,
 )
 
 _T = TypeVar('_T')
@@ -82,6 +82,7 @@ def __lt__(self, n: int) -> bool: pass
     def __gt__(self, n: int) -> bool: pass
     def __le__(self, n: int) -> bool: pass
     def __ge__(self, n: int) -> bool: pass
+    def to_bytes(self, length: int, order: str, signed: bool = False) -> bytes: pass
 
 class str:
     @overload

mypyc/test-data/irbuild-int.test

@@ -210,3 +210,26 @@ L0:
     r0 = CPyTagged_Invert(n)
     x = r0
     return x
+
+[case testIntToBytes]
+def f(x: int) -> bytes:
+    return x.to_bytes(2, "big")
+def g(x: int) -> bytes:
+    return x.to_bytes(4, "little", True)
+[out]
+def f(x):
+    x :: int
+    r0 :: str
+    r1 :: bytes
+L0:
+    r0 = 'big'
+    r1 = CPyTagged_ToBytes(x, 4, r0, 0)
+    return r1
+def g(x):
+    x :: int
+    r0 :: str
+    r1 :: bytes
+L0:
+    r0 = 'little'
+    r1 = CPyTagged_ToBytes(x, 8, r0, 1)
+    return r1

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'