Add long_asnativebytes() function

vstinner · vstinner · commit 5d3e22496011 · 2024-09-17T13:10:30.000+02:00
diff --git a/Objects/longobject.c b/Objects/longobject.c
@@ -1117,53 +1117,18 @@ _resolve_endianness(int *endianness)
     return 0;
 }
 
-Py_ssize_t
-PyLong_AsNativeBytes(PyObject* vv, void* buffer, Py_ssize_t n, int flags)
+
+static Py_ssize_t
+long_asnativebytes(PyLongObject *v, void* buffer, Py_ssize_t n,
+                   int little_endian, int unsigned_buffer)
 {
-    PyLongObject *v;
+    Py_ssize_t res = 0;
     union {
         Py_ssize_t v;
         unsigned char b[sizeof(Py_ssize_t)];
     } cv;
-    int do_decref = 0;
-    Py_ssize_t res = 0;
-
-    if (vv == NULL || n < 0) {
-        PyErr_BadInternalCall();
-        return -1;
-    }
-
-    int little_endian = flags;
-    if (_resolve_endianness(&little_endian) < 0) {
-        return -1;
-    }
-
-    if (PyLong_Check(vv)) {
-        v = (PyLongObject *)vv;
-    }
-    else if (flags != -1 && (flags & Py_ASNATIVEBYTES_ALLOW_INDEX)) {
-        v = (PyLongObject *)_PyNumber_Index(vv);
-        if (v == NULL) {
-            return -1;
-        }
-        do_decref = 1;
-    }
-    else {
-        PyErr_Format(PyExc_TypeError, "expect int, got %T", vv);
-        return -1;
-    }
-
-    if ((flags != -1 && (flags & Py_ASNATIVEBYTES_REJECT_NEGATIVE))
-        && _PyLong_IsNegative(v)) {
-        PyErr_SetString(PyExc_ValueError, "Cannot convert negative int");
-        if (do_decref) {
-            Py_DECREF(v);
-        }
-        return -1;
-    }
 
     if (_PyLong_IsCompact(v)) {
-        res = 0;
         cv.v = _PyLong_CompactValue(v);
         /* Most paths result in res = sizeof(compact value). Only the case
          * where 0 < n < sizeof(compact value) do we need to check and adjust
@@ -1200,7 +1165,7 @@ PyLong_AsNativeBytes(PyObject* vv, void* buffer, Py_ssize_t n, int flags)
                 /* Positive values with the MSB set do not require an
                  * additional bit when the caller's intent is to treat them
                  * as unsigned. */
-                if (flags == -1 || (flags & Py_ASNATIVEBYTES_UNSIGNED_BUFFER)) {
+                if (unsigned_buffer) {
                     res = n;
                 } else {
                     res = n + 1;
@@ -1287,7 +1252,7 @@ PyLong_AsNativeBytes(PyObject* vv, void* buffer, Py_ssize_t n, int flags)
                  * as unsigned. */
                 unsigned char *b = (unsigned char *)buffer;
                 if (b[little_endian ? n - 1 : 0] & 0x80) {
-                    if (flags == -1 || (flags & Py_ASNATIVEBYTES_UNSIGNED_BUFFER)) {
+                    if (unsigned_buffer) {
                         res = n;
                     } else {
                         res = n + 1;
@@ -1296,6 +1261,54 @@ PyLong_AsNativeBytes(PyObject* vv, void* buffer, Py_ssize_t n, int flags)
             }
         }
     }
+    return res;
+}
+
+
+Py_ssize_t
+PyLong_AsNativeBytes(PyObject* vv, void* buffer, Py_ssize_t n, int flags)
+{
+    PyLongObject *v;
+    int do_decref = 0;
+    Py_ssize_t res = 0;
+
+    if (vv == NULL || n < 0) {
+        PyErr_BadInternalCall();
+        return -1;
+    }
+
+    int little_endian = flags;
+    if (_resolve_endianness(&little_endian) < 0) {
+        return -1;
+    }
+
+    if (PyLong_Check(vv)) {
+        v = (PyLongObject *)vv;
+    }
+    else if (flags != -1 && (flags & Py_ASNATIVEBYTES_ALLOW_INDEX)) {
+        v = (PyLongObject *)_PyNumber_Index(vv);
+        if (v == NULL) {
+            return -1;
+        }
+        do_decref = 1;
+    }
+    else {
+        PyErr_Format(PyExc_TypeError, "expect int, got %T", vv);
+        return -1;
+    }
+
+    if ((flags != -1 && (flags & Py_ASNATIVEBYTES_REJECT_NEGATIVE))
+        && _PyLong_IsNegative(v)) {
+        PyErr_SetString(PyExc_ValueError, "Cannot convert negative int");
+        if (do_decref) {
+            Py_DECREF(v);
+        }
+        return -1;
+    }
+
+    int unsigned_buffer = (flags == -1
+                           || (flags & Py_ASNATIVEBYTES_UNSIGNED_BUFFER));
+    res = long_asnativebytes(v, buffer, n, little_endian, unsigned_buffer);
 
     if (do_decref) {
         Py_DECREF(v);
@@ -6797,13 +6810,11 @@ PyLong_Export(PyObject *obj, PyLongExport *export_long)
         PyErr_Format(PyExc_TypeError, "expect int, got %T", obj);
         return -1;
     }
+    PyLongObject *self = (PyLongObject*)obj;
 
     int64_t value;
-    int flags = Py_ASNATIVEBYTES_NATIVE_ENDIAN;
-    Py_ssize_t bytes = PyLong_AsNativeBytes(obj, &value, sizeof(value), flags);
-    if (bytes < 0) {
-        return -1;
-    }
+    Py_ssize_t bytes = long_asnativebytes(self, &value, sizeof(value),
+                                          PY_LITTLE_ENDIAN, 0);
 
     if ((size_t)bytes <= sizeof(value)) {
         export_long->value = value;
@@ -6813,7 +6824,6 @@ PyLong_Export(PyObject *obj, PyLongExport *export_long)
         export_long->_reserved = 0;
     }
     else {
-        PyLongObject *self = (PyLongObject*)obj;
         export_long->value = 0;
         export_long->negative = _PyLong_IsNegative(self);
         export_long->ndigits = _PyLong_DigitCount(self);
