Add FromBits functions

python · zooba · Feb 12, 2024 · Feb 2, 2024 · Feb 2, 2024 · Feb 2, 2024
commit 79d4942fed914342639b43b44b4c40aeac431105
diff --git a/Include/cpython/longobject.h b/Include/cpython/longobject.h
@@ -22,17 +22,17 @@ PyAPI_FUNC(PyObject*) PyLong_FromUnicodeObject(PyObject *u, int base);
 PyAPI_FUNC(int) PyLong_CopyBits(PyObject* v, void* buffer, size_t n_bytes,
     int endianness);
 
-/* PyLong_FromByteArray: Create an integer value containing the number from
-   a native buffer.
-   n is the number of bytes to read from the buffer.
-   Uses the current build's default endianness, and assumes the value was
-   sign extended to 'n' bytes.
-   PyLong_FromUnsignedByteArray assumes the value was zero extended, and
-   even if the MSB is set the resulting int will be positive.
+/* PyLong_FromBits: Create an int value from a native integer
+   n_bytes is the number of bytes to read from the buffer. Passing 0 will
+   always produce the zero int.
+   PyLong_FromUnsignedBits always produces a positive int.
+   endianness is -1 for native endian, 0 for big endian or 1 for little.
 
    Returns the int object, or NULL with an exception set. */
-PyAPI_FUNC(PyObject) PyLong_FromByteArray(void* buffer, size_t n);
-PyAPI_FUNC(PyObject) PyLong_FromUnsignedByteArray(void* buffer, size_t n);
+PyAPI_FUNC(PyObject*) PyLong_FromBits(void* buffer, size_t n_bytes,
+    int endianness);
+PyAPI_FUNC(PyObject*) PyLong_FromUnsignedBits(void* buffer, size_t n_bytes,
+    int endianness);
 
 PyAPI_FUNC(int) PyUnstable_Long_IsCompact(const PyLongObject* op);
 PyAPI_FUNC(Py_ssize_t) PyUnstable_Long_CompactValue(const PyLongObject* op);

diff --git a/Lib/test/test_capi/test_long.py b/Lib/test/test_capi/test_long.py
@@ -426,16 +426,10 @@ def test_long_asvoidptr(self):
 
     def test_long_copybits(self):
         import math
-        from _testcapi import (
-            pylong_copybits as copybits,
-            SIZE_MAX,
-        )
+        from _testcapi import pylong_copybits as copybits, SIZE_MAX
 
-        def log2(x):
-            return math.log(x) / math.log(2)
-
-        # Abbreviated sizeof(Py_ssize_t) because we use it a lot
-        SZ = int(math.ceil(log2(SIZE_MAX + 1)) / 8)
+        # Abbreviate sizeof(Py_ssize_t) to SZ because we use it a lot
+        SZ = int(math.ceil(math.log(SIZE_MAX + 1) / math.log(2)) / 8)
         MAX_SSIZE = 2 ** (SZ * 8 - 1) - 1
         MAX_USIZE = 2 ** (SZ * 8) - 1
         if support.verbose:
@@ -520,5 +514,34 @@ def log2(x):
                     f"PyLong_CopyBits(v, buffer, {n}, <little>)")
                 self.assertEqual(expect_le, buffer[:n], "<little>")
 
+    def test_long_frombits(self):
+        import math
+        from _testcapi import pylong_frombits as frombits, SIZE_MAX
+
+        # Abbreviate sizeof(Py_ssize_t) to SZ because we use it a lot
+        SZ = int(math.ceil(math.log(SIZE_MAX + 1) / math.log(2)) / 8)
+        MAX_SSIZE = 2 ** (SZ * 8 - 1) - 1
+        MAX_USIZE = 2 ** (SZ * 8) - 1
+
+        for v_be, expect_s, expect_u in [
+            (b'\x00', 0, 0),
+            (b'\x01', 1, 1),
+            (b'\xff', -1, 255),
+            (b'\x00\xff', 255, 255),
+            (b'\xff\xff', -1, 65535),
+        ]:
+            with self.subTest(f"{expect_s}-{expect_u:X}-{len(v_be)}bytes"):
+                n = len(v_be)
+                v_le = v_be[::-1]
+
+                self.assertEqual(expect_s, frombits(v_be, n, 0, 1),
+                    f"PyLong_FromBits(buffer, {n}, <big>)")
+                self.assertEqual(expect_s, frombits(v_le, n, 1, 1),
+                    f"PyLong_FromBits(buffer, {n}, <little>)")
+                self.assertEqual(expect_u, frombits(v_be, n, 0, 0),
+                    f"PyLong_FromUnsignedBits(buffer, {n}, <big>)")
+                self.assertEqual(expect_u, frombits(v_le, n, 1, 0),
+                    f"PyLong_FromUnsignedBits(buffer, {n}, <little>)")
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Modules/_testcapi/long.c b/Modules/_testcapi/long.c
@@ -800,6 +800,27 @@ pylong_copybits(PyObject *module, PyObject *args)
     return res >= 0 ? PyLong_FromLong(res) : NULL;
 }
 
+static PyObject *
+pylong_frombits(PyObject *module, PyObject *args)
+{
+    Py_buffer buffer;
+    Py_ssize_t n, endianness, signed_;
+    if (!PyArg_ParseTuple(args, "y*nnn", &buffer, &n, &endianness, &signed_)) {
+        return NULL;
+    }
+    if (buffer.len < n) {
+        PyErr_SetString(PyExc_ValueError, "buffer must be at least 'n' bytes");
+        PyBuffer_Release(&buffer);
+        return NULL;
+    }
+    PyObject *res = signed_ 
+        ? PyLong_FromBits(buffer.buf, n, (int)endianness)
+        : PyLong_FromUnsignedBits(buffer.buf, n, (int)endianness);
+    PyBuffer_Release(&buffer);
+    return res;
+}
+
+
 static PyMethodDef test_methods[] = {
     _TESTCAPI_TEST_LONG_AND_OVERFLOW_METHODDEF
     _TESTCAPI_TEST_LONG_API_METHODDEF
@@ -829,6 +850,7 @@ static PyMethodDef test_methods[] = {
     {"pylong_asdouble",             pylong_asdouble,            METH_O},
     {"pylong_asvoidptr",            pylong_asvoidptr,           METH_O},
     {"pylong_copybits",             pylong_copybits,            METH_VARARGS},
+    {"pylong_frombits",             pylong_frombits,            METH_VARARGS},
     {NULL},
 };
 

diff --git a/Objects/longobject.c b/Objects/longobject.c
@@ -1217,6 +1217,50 @@ PyLong_CopyBits(PyObject* vv, void* buffer, size_t n, int endianness)
 }
 
 
+PyObject *
+PyLong_FromBits(const void* buffer, size_t n, int endianness)
+{
+    if (!buffer) {
+        PyErr_BadInternalCall();
+        return NULL;
+    }
+
+    int little_endian = endianness;
+    if (endianness < 0) {
+        endianness = PY_LITTLE_ENDIAN;
+    }
+    if (endianness != 0 && endianness != 1) {
+        PyErr_SetString(PyExc_SystemError, "invalid 'endianness' value");
+        return NULL;
+    }
+
+    return _PyLong_FromByteArray((const unsigned char *)buffer, n,
+                                 little_endian, 1);
+}
+
+
+PyObject *
+PyLong_FromUnsignedBits(void* buffer, size_t n, int endianness)
+{
+    if (!buffer) {
+        PyErr_BadInternalCall();
+        return NULL;
+    }
+
+    int little_endian = endianness;
+    if (endianness < 0) {
+        endianness = PY_LITTLE_ENDIAN;
+    }
+    if (endianness != 0 && endianness != 1) {
+        PyErr_SetString(PyExc_SystemError, "invalid 'endianness' value");
+        return NULL;
+    }
+
+    return _PyLong_FromByteArray((const unsigned char *)buffer, n,
+                                 little_endian, 0);
+}
+
+
 /* Create a new int object from a C pointer */
 
 PyObject *