gh-102471, PEP 757: Add PyLong import and export API (#121339)

Co-authored-by: Sergey B Kirpichev <skirpichev@gmail.com>
Co-authored-by: Steve Dower <steve.dower@microsoft.com>
Co-authored-by: Bénédikt Tran <10796600+picnixz@users.noreply.github.com>

1 files changed, 91 insertions, 0 deletions

diff --git a/Lib/test/test_capi/test_long.py b/Lib/test/test_capi/test_long.py
index a770945..d45ac75 100644
--- a/Lib/test/test_capi/test_long.py
+++ b/Lib/test/test_capi/test_long.py
@@ -10,6 +10,7 @@ _testlimitedcapi = import_helper.import_module('_testlimitedcapi')
 
 NULL = None
 
+
 class IntSubclass(int):
     pass
 
@@ -714,5 +715,95 @@ class LongTests(unittest.TestCase):
         self.check_long_asint(as_uint64, 0, UINT64_MAX,
                               negative_value_error=ValueError)
 
+    def test_long_layout(self):
+        # Test PyLong_GetNativeLayout()
+        int_info = sys.int_info
+        layout = _testcapi.get_pylong_layout()
+        expected = {
+            'bits_per_digit': int_info.bits_per_digit,
+            'digit_size': int_info.sizeof_digit,
+            'digits_order': -1,
+            'digit_endianness': -1 if sys.byteorder == 'little' else 1,
+        }
+        self.assertEqual(layout, expected)
+
+    def test_long_export(self):
+        # Test PyLong_Export()
+        layout = _testcapi.get_pylong_layout()
+        base = 2 ** layout['bits_per_digit']
+
+        pylong_export = _testcapi.pylong_export
+
+        # value fits into int64_t
+        self.assertEqual(pylong_export(0), 0)
+        self.assertEqual(pylong_export(123), 123)
+        self.assertEqual(pylong_export(-123), -123)
+        self.assertEqual(pylong_export(IntSubclass(123)), 123)
+
+        # use an array, doesn't fit into int64_t
+        self.assertEqual(pylong_export(base**10 * 2 + 1),
+                         (0, [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2]))
+        self.assertEqual(pylong_export(-(base**10 * 2 + 1)),
+                         (1, [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2]))
+        self.assertEqual(pylong_export(IntSubclass(base**10 * 2 + 1)),
+                         (0, [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2]))
+
+        self.assertRaises(TypeError, pylong_export, 1.0)
+        self.assertRaises(TypeError, pylong_export, 0+1j)
+        self.assertRaises(TypeError, pylong_export, "abc")
+
+    def test_longwriter_create(self):
+        # Test PyLongWriter_Create()
+        layout = _testcapi.get_pylong_layout()
+        base = 2 ** layout['bits_per_digit']
+
+        pylongwriter_create = _testcapi.pylongwriter_create
+        self.assertRaises(ValueError, pylongwriter_create, 0, [])
+        self.assertRaises(ValueError, pylongwriter_create, -123, [])
+        self.assertEqual(pylongwriter_create(0, [0]), 0)
+        self.assertEqual(pylongwriter_create(0, [123]), 123)
+        self.assertEqual(pylongwriter_create(1, [123]), -123)
+        self.assertEqual(pylongwriter_create(1, [1, 2]),
+                         -(base * 2 + 1))
+        self.assertEqual(pylongwriter_create(0, [1, 2, 3]),
+                         base**2 * 3 + base * 2 + 1)
+        max_digit = base - 1
+        self.assertEqual(pylongwriter_create(0, [max_digit, max_digit, max_digit]),
+                         base**2 * max_digit + base * max_digit + max_digit)
+
+        # normalize
+        self.assertEqual(pylongwriter_create(0, [123, 0, 0]), 123)
+
+        # test singletons + normalize
+        for num in (-2, 0, 1, 5, 42, 100):
+            self.assertIs(pylongwriter_create(bool(num < 0), [abs(num), 0]),
+                          num)
+
+        def to_digits(num):
+            digits = []
+            while True:
+                num, digit = divmod(num, base)
+                digits.append(digit)
+                if not num:
+                    break
+            return digits
+
+        # round trip: Python int -> export -> Python int
+        pylong_export = _testcapi.pylong_export
+        numbers = [*range(0, 10), 12345, 0xdeadbeef, 2**100, 2**100-1]
+        numbers.extend(-num for num in list(numbers))
+        for num in numbers:
+            with self.subTest(num=num):
+                data = pylong_export(num)
+                if isinstance(data, tuple):
+                    negative, digits = data
+                else:
+                    value = data
+                    negative = int(value < 0)
+                    digits = to_digits(abs(value))
+                self.assertEqual(pylongwriter_create(negative, digits), num,
+                                 (negative, digits))
+
+
 if __name__ == "__main__":
     unittest.main()