Better options handling and tests

python · zooba · Feb 12, 2024 · Feb 2, 2024 · Feb 2, 2024 · Feb 2, 2024
commit bf86f851fe4e2f05efeffb94d786be4b26fb8d44
diff --git a/Include/cpython/longobject.h b/Include/cpython/longobject.h
@@ -24,19 +24,22 @@ PyAPI_FUNC(int) PyLong_AsUnsignedByteArray(PyObject* v, void* buffer, size_t n);
 PyAPI_FUNC(int) PyLong_AsByteArrayWithOptions(PyObject* v, void* buffer,
                                               size_t n, int options);
 
-#define PYLONG_ASBYTEARRAY_NATIVE_ENDIAN    0x00
-#define PYLONG_ASBYTEARRAY_LITTLE_ENDIAN    0x01
-#define PYLONG_ASBYTEARRAY_BIG_ENDIAN       0x02
+/* Deliberately avoiding values 0 and 1 to avoid letting people
+   accidentally pass PY_LITTLE_ENDIAN as a constant. */
+#define PYLONG_ASBYTEARRAY_LITTLE_ENDIAN    0x04
+#define PYLONG_ASBYTEARRAY_BIG_ENDIAN       0x08
+#define PYLONG_ASBYTEARRAY_NATIVE_ENDIAN    (0x04|0x08)
 
-#define PYLONG_ASBYTEARRAY_SIGNED   0x00
-#define PYLONG_ASBYTEARRAY_UNSIGNED 0x04
+#define PYLONG_ASBYTEARRAY_SIGNED   0x10
+#define PYLONG_ASBYTEARRAY_UNSIGNED 0x20
 
 /* 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.
+   PyLong_FromUnsignedByteArray assumes the value was zero extended, and
+   even if the MSB is set the resulting int will be positive.
 
    Returns the int object, or NULL with an exception set. */
 PyAPI_FUNC(PyObject) PyLong_FromByteArray(void* buffer, size_t n);

diff --git a/Lib/test/test_capi/test_long.py b/Lib/test/test_capi/test_long.py
@@ -430,7 +430,12 @@ def test_long_asbytearray(self):
             pylong_asbytearray as asbytearray,
             pylong_asunsignedbytearray as asunsignedbytearray,
             pylong_asbytearraywithoptions as asbytearraywithoptions,
-            SIZE_MAX
+            SIZE_MAX,
+            PYLONG_ASBYTEARRAY_NATIVE_ENDIAN,
+            PYLONG_ASBYTEARRAY_LITTLE_ENDIAN,
+            PYLONG_ASBYTEARRAY_BIG_ENDIAN,
+            PYLONG_ASBYTEARRAY_SIGNED,
+            PYLONG_ASBYTEARRAY_UNSIGNED,
         )
 
         def log2(x):
@@ -442,6 +447,19 @@ def log2(x):
         if support.verbose:
             print(f"{SIZEOF_SIZE=}\n{MAX_SSIZE=:016X}\n{MAX_USIZE=:016X}")
 
+        U_LE = PYLONG_ASBYTEARRAY_UNSIGNED | PYLONG_ASBYTEARRAY_LITTLE_ENDIAN
+        U_BE = PYLONG_ASBYTEARRAY_UNSIGNED | PYLONG_ASBYTEARRAY_BIG_ENDIAN
+        S_LE = PYLONG_ASBYTEARRAY_SIGNED | PYLONG_ASBYTEARRAY_LITTLE_ENDIAN
+        S_BE = PYLONG_ASBYTEARRAY_SIGNED | PYLONG_ASBYTEARRAY_BIG_ENDIAN
+
+        # Ensure options of 0 and 1 raise
+        with self.assertRaises(SystemError):
+            asbytearraywithoptions(0, bytearray(1), 0, 0)
+        with self.assertRaises(SystemError):
+            asbytearraywithoptions(0, bytearray(1), 0, 1)
+
+        # These tests request the required buffer size for both unsigned and
+        # signed options.
         for v, expect_u, expect_s in [
             (0, SIZEOF_SIZE, SIZEOF_SIZE),
             (512, SIZEOF_SIZE, SIZEOF_SIZE),
@@ -450,22 +468,28 @@ def log2(x):
             (MAX_USIZE, SIZEOF_SIZE, SIZEOF_SIZE + 1),
             (-MAX_SSIZE, SIZEOF_SIZE, SIZEOF_SIZE),
             (-MAX_USIZE, SIZEOF_SIZE, SIZEOF_SIZE + 1),
+            (2**255-1, 32, 32),
+            (-(2**255-1), 32, 32),
+            (2**256-1, 32, 33),
+            (-(2**256-1), 32, 33),
         ]:
             with self.subTest(f"sizeof-{v:X}"):
                 buffer = bytearray(1)
                 self.assertEqual(expect_u, asunsignedbytearray(v, buffer, 0),
                     "PyLong_AsUnsignedByteArray(v, NULL, 0)")
                 self.assertEqual(expect_s, asbytearray(v, buffer, 0),
                     "PyLong_AsByteArray(v, NULL, 0)")
-                self.assertEqual(expect_u, asbytearraywithoptions(v, buffer, 0, 0x05),
+                self.assertEqual(expect_u, asbytearraywithoptions(v, buffer, 0, U_LE),
                     "PyLong_AsByteArrayWithOptions(v, NULL, 0, unsigned|little)")
-                self.assertEqual(expect_u, asbytearraywithoptions(v, buffer, 0, 0x06),
+                self.assertEqual(expect_u, asbytearraywithoptions(v, buffer, 0, U_BE),
                     "PyLong_AsByteArrayWithOptions(v, NULL, 0, unsigned|big)")
-                self.assertEqual(expect_s, asbytearraywithoptions(v, buffer, 0, 0x01),
+                self.assertEqual(expect_s, asbytearraywithoptions(v, buffer, 0, S_LE),
                     "PyLong_AsByteArrayWithOptions(v, NULL, 0, signed|little)")
-                self.assertEqual(expect_s, asbytearraywithoptions(v, buffer, 0, 0x02),
+                self.assertEqual(expect_s, asbytearraywithoptions(v, buffer, 0, S_BE),
                     "PyLong_AsByteArrayWithOptions(v, NULL, 0, signed|big)")
 
+        # We request as many bytes as `expect_be` contains, so tests for the
+        # same value may test different things with a longer expected result.
         for v, expect_be, signed_fails in [
             (0, b'\x00', False),
             (0, b'\x00' * 2, False),
@@ -475,6 +499,7 @@ def log2(x):
             (42, b'\x2a', False),
             (42, b'\x00' * 10 + b'\x2a', False),
             (-1, b'\xff', False),
+            (-1, b'\xff' * 10, False),
             (-42, b'\xd6', False),
             (-42, b'\xff' * 10 + b'\xd6', False),
             # Only unsigned will extract 255 into a single byte
@@ -484,36 +509,50 @@ def log2(x):
             (2**63, b'\x80\x00\x00\x00\x00\x00\x00\x00', True),
             (-2**63, b'\x80\x00\x00\x00\x00\x00\x00\x00', False),
             (2**63, b'\x00\x80\x00\x00\x00\x00\x00\x00\x00', False),
-            (-2**63, b'\xFF\x80\x00\x00\x00\x00\x00\x00\x00', False),
+            (-2**63, b'\xff\x80\x00\x00\x00\x00\x00\x00\x00', False),
+            (2**255-1, b'\x7f' + b'\xff' * 31, False),
+            (-(2**255-1), b'\x80' + b'\x00' * 30 + b'\x01', False),
+            (-(2**255-1), b'\xff\x80' + b'\x00' * 30 + b'\x01', False),
+            # Cannot extract 256 bits into 32 bytes as signed ...
+            (2**256-1, b'\xff' * 32, True),
+            # ... but can extract into 33 bytes (zero extended)
+            (2**256-1, b'\x00' + b'\xff' * 32, False),
+
+            # Unsigned extract of this negative number to 32 bytes will succeed,
+            # sacrificing the sign bit and overall magnitude, just like a C
+            # cast (int256_t)0x1_0000_..._0001 would. But signed extract fails
+            (-(2**256-1), b'\x00' * 31 + b'\x01', True),
+            # Extract to 33 bytes preserves the top-most bits in all cases.
+            (-(2**256-1), b'\xff' + b'\x00' * 31 + b'\x01', False),
         ]:
             with self.subTest(f"{v:X}-{len(expect_be)}bytes"):
                 n = len(expect_be)
                 buffer = bytearray(n)
                 expect_le = expect_be[::-1]
 
-                self.assertEqual(0, asbytearraywithoptions(v, buffer, n, 0x05),
+                self.assertEqual(0, asbytearraywithoptions(v, buffer, n, U_LE),
                     f"PyLong_AsByteArrayWithOptions(v, buffer, {n}, unsigned+little)")
                 self.assertEqual(expect_le, buffer[:n], "unsigned+little")
-                self.assertEqual(0, asbytearraywithoptions(v, buffer, n, 0x06),
+                self.assertEqual(0, asbytearraywithoptions(v, buffer, n, U_BE),
                     f"PyLong_AsByteArrayWithOptions(v, buffer, {n}, unsigned+big)")
                 self.assertEqual(expect_be, buffer[:n], "unsigned+big")
 
                 if signed_fails:
                     self.assertEqual(
                         max(n + 1, SIZEOF_SIZE),
-                        asbytearraywithoptions(v, buffer, n, 0x01),
+                        asbytearraywithoptions(v, buffer, n, S_LE),
                         f"PyLong_AsByteArrayWithOptions(v, buffer, {n}, signed+little)",
                     )
                     self.assertEqual(
                         max(n + 1, SIZEOF_SIZE),
-                        asbytearraywithoptions(v, buffer, n, 0x02),
+                        asbytearraywithoptions(v, buffer, n, S_BE),
                         f"PyLong_AsByteArrayWithOptions(v, buffer, {n}, signed+big)",
                     )
                 else:
-                    self.assertEqual(0, asbytearraywithoptions(v, buffer, n, 0x01),
+                    self.assertEqual(0, asbytearraywithoptions(v, buffer, n, S_LE),
                         f"PyLong_AsByteArrayWithOptions(v, buffer, {n}, signed+little)")
                     self.assertEqual(expect_le, buffer[:n], "signed+little")
-                    self.assertEqual(0, asbytearraywithoptions(v, buffer, n, 0x02),
+                    self.assertEqual(0, asbytearraywithoptions(v, buffer, n, S_BE),
                         f"PyLong_AsByteArrayWithOptions(v, buffer, {n}, signed+big)")
                     self.assertEqual(expect_be, buffer[:n], "signed+big")
 

diff --git a/Modules/_testcapi/long.c b/Modules/_testcapi/long.c
@@ -888,6 +888,10 @@ _PyTestCapi_Init_Long(PyObject *mod)
     if (PyModule_AddFunctions(mod, test_methods) < 0) {
         return -1;
     }
-
+    if (PyModule_AddIntMacro(mod, PYLONG_ASBYTEARRAY_NATIVE_ENDIAN)) return -1;
+    if (PyModule_AddIntMacro(mod, PYLONG_ASBYTEARRAY_LITTLE_ENDIAN)) return -1;
+    if (PyModule_AddIntMacro(mod, PYLONG_ASBYTEARRAY_BIG_ENDIAN)) return -1;
+    if (PyModule_AddIntMacro(mod, PYLONG_ASBYTEARRAY_SIGNED)) return -1;
+    if (PyModule_AddIntMacro(mod, PYLONG_ASBYTEARRAY_UNSIGNED)) return -1;
     return 0;
 }
diff --git a/Objects/longobject.c b/Objects/longobject.c
@@ -1075,8 +1075,11 @@ _fits_in_n_bits(Py_ssize_t v, Py_ssize_t n, int require_sign_bit)
     return v_extended == 0 || v_extended == -1;
 }
 
-int
-PyLong_AsByteArrayWithOptions(PyObject* vv, void* buffer, size_t n, int options)
+// Private function means we can add more options safely.
+// The public API parses a flags parameter to extract the requested settings.
+static int
+_PyLong_AsByteArrayWithOptions(PyObject* vv, void* buffer, size_t n,
+    int signed_, int little_endian)
 {
     PyLongObject *v;
     union {
@@ -1085,21 +1088,6 @@ PyLong_AsByteArrayWithOptions(PyObject* vv, void* buffer, size_t n, int options)
     } cv;
     int do_decref = 0;
     int res = 0;
-    int signed_ = !(options & PYLONG_ASBYTEARRAY_UNSIGNED);
-    int little_endian = PY_LITTLE_ENDIAN;
-    switch (options & (PYLONG_ASBYTEARRAY_LITTLE_ENDIAN | PYLONG_ASBYTEARRAY_BIG_ENDIAN)) {
-    case 0:
-        break;
-    case PYLONG_ASBYTEARRAY_BIG_ENDIAN:
-        little_endian = 0;
-        break;
-    case PYLONG_ASBYTEARRAY_LITTLE_ENDIAN:
-        little_endian = 1;
-        break;
-    default:
-        PyErr_SetString(PyExc_ValueError, "invalid 'options' value");
-        return -1;
-    }
 
     if (vv == NULL) {
         PyErr_BadInternalCall();
@@ -1191,12 +1179,12 @@ PyLong_AsByteArrayWithOptions(PyObject* vv, void* buffer, size_t n, int options)
         res = -1;
     }
     else if (!signed_ && _PyLong_IsNegative(v)) {
-        // This case is not supported by _PyLong_AsByteArray, so we extract a
-        // signed value, to a larger buffer first if needed. If we use a larger
-        // buffer, the extra data should be all bits set. If it is, then we can
-        // ignore it and return the number of bytes the caller requested
-        // (that's what the "unsigned" means). Otherwise, we need to return the
-        // actual number of bytes required.
+        /* This case is not supported by _PyLong_AsByteArray, so we extract a
+         * signed value, to a larger buffer first if needed. If we use a larger
+         * buffer, the extra data should be all bits set. If it is, then we can
+         * ignore it and return the number of bytes the caller requested
+         * (that's what the "unsigned" means). Otherwise, we need to return the
+         * actual number of bytes required. */
         res = _PyLong_AsByteArray(v, buffer, n, little_endian, 1, 0);
         if (res < 0) {
             unsigned char *b = (unsigned char *)PyMem_Malloc(n + 1);
@@ -1208,7 +1196,8 @@ PyLong_AsByteArrayWithOptions(PyObject* vv, void* buffer, size_t n, int options)
             if (res == 0) {
                 // Ensure the extra byte is 0xFF
                 if (little_endian) {
-                    if (b[n - 1] != 0xFF) {
+                    // (remember, we allocated n+1 bytes)
+                    if (b[n] != 0xFF) {
                         res = -1;
                     } else {
                         memcpy(buffer, b, n);
@@ -1228,9 +1217,9 @@ PyLong_AsByteArrayWithOptions(PyObject* vv, void* buffer, size_t n, int options)
         res = _PyLong_AsByteArray(v, buffer, n, little_endian, signed_, 0);
     }
 
-    // NOTE: res should only be <0 if a _PyLong_AsByteArray has failed without
-    // setting an exception, or we're requesting the size of a non-compact int.
-    // If you add another reason, you should update this!
+    /* NOTE: res should only be <0 if a _PyLong_AsByteArray has failed without
+     * setting an exception, or we're requesting the size of a non-compact int.
+     * If you add another reason, you should update this! */
     if (res < 0) {
         // More efficient calculation for number of bytes required?
         size_t nb = _PyLong_NumBits((PyObject *)v) + (signed_ ? 1 : 0);
@@ -1251,16 +1240,64 @@ PyLong_AsByteArrayWithOptions(PyObject* vv, void* buffer, size_t n, int options)
     return res;
 }
 
+int
+PyLong_AsByteArrayWithOptions(PyObject* vv, void* buffer, size_t n, int options)
+{
+    int signed_;
+    int little_endian;
+
+    // option 1 is excluded and always raises
+    if (!options || options & 0x01) {
+        PyErr_SetString(PyExc_SystemError, "invalid 'options' value");
+        return -1;
+    }
+
+    switch (options & (PYLONG_ASBYTEARRAY_SIGNED | PYLONG_ASBYTEARRAY_UNSIGNED)) {
+    case PYLONG_ASBYTEARRAY_SIGNED:
+        signed_ = 1;
+        break;
+    case PYLONG_ASBYTEARRAY_UNSIGNED:
+        signed_ = 0;
+        break;
+    case 0:
+        PyErr_SetString(PyExc_SystemError, "invalid 'options' value - no sign option");
+        return -1;
+    default:
+        PyErr_SetString(PyExc_SystemError, "invalid 'options' value");
+        return -1;
+    }
+
+    switch (options & (PYLONG_ASBYTEARRAY_LITTLE_ENDIAN | PYLONG_ASBYTEARRAY_BIG_ENDIAN
+                       | PYLONG_ASBYTEARRAY_NATIVE_ENDIAN)) {
+    case PYLONG_ASBYTEARRAY_BIG_ENDIAN:
+        little_endian = 0;
+        break;
+    case PYLONG_ASBYTEARRAY_LITTLE_ENDIAN:
+        little_endian = 1;
+        break;
+    case PYLONG_ASBYTEARRAY_NATIVE_ENDIAN:
+        little_endian = PY_LITTLE_ENDIAN;
+        break;
+    case 0:
+        PyErr_SetString(PyExc_SystemError, "invalid 'options' value - no endian option");
+        return -1;
+    default:
+        PyErr_SetString(PyExc_SystemError, "invalid 'options' value");
+        return -1;
+    }
+    return _PyLong_AsByteArrayWithOptions(vv, buffer, n, signed_, little_endian);
+}
+
 int
 PyLong_AsByteArray(PyObject* vv, void* buffer, size_t n)
 {
-    return PyLong_AsByteArrayWithOptions(vv, buffer, n, PYLONG_ASBYTEARRAY_SIGNED);
+    return _PyLong_AsByteArrayWithOptions(vv, buffer, n, 1, PY_LITTLE_ENDIAN);
 }
 
 int
 PyLong_AsUnsignedByteArray(PyObject* vv, void* buffer, size_t n)
 {
-    return PyLong_AsByteArrayWithOptions(vv, buffer, n, PYLONG_ASBYTEARRAY_UNSIGNED);
+    return _PyLong_AsByteArrayWithOptions(vv, buffer, n, 0, PY_LITTLE_ENDIAN);
 }