gh-121249: Support _Complex types in the struct module (#121613)

Co-authored-by: Peter Bierma <zintensitydev@gmail.com>
Co-authored-by: Bénédikt Tran <10796600+picnixz@users.noreply.github.com>
Co-authored-by: Victor Stinner <vstinner@python.org>

1 files changed, 36 insertions, 1 deletions

diff --git a/Lib/test/test_struct.py b/Lib/test/test_struct.py
index bdbf880..e3193c7 100644
--- a/Lib/test/test_struct.py
+++ b/Lib/test/test_struct.py
@@ -1,4 +1,5 @@
 from collections import abc
+from itertools import combinations
 import array
 import gc
 import math
@@ -11,12 +12,22 @@ import weakref
 from test import support
 from test.support import import_helper, suppress_immortalization
 from test.support.script_helper import assert_python_ok
+from test.support.testcase import ComplexesAreIdenticalMixin
 
 ISBIGENDIAN = sys.byteorder == "big"
 
 integer_codes = 'b', 'B', 'h', 'H', 'i', 'I', 'l', 'L', 'q', 'Q', 'n', 'N'
 byteorders = '', '@', '=', '<', '>', '!'
 
+INF = float('inf')
+NAN = float('nan')
+
+try:
+    struct.pack('C', 1j)
+    have_c_complex = True
+except struct.error:
+    have_c_complex = False
+
 def iter_integer_formats(byteorders=byteorders):
     for code in integer_codes:
         for byteorder in byteorders:
@@ -33,7 +44,7 @@ def bigendian_to_native(value):
     else:
         return string_reverse(value)
 
-class StructTest(unittest.TestCase):
+class StructTest(ComplexesAreIdenticalMixin, unittest.TestCase):
     def test_isbigendian(self):
         self.assertEqual((struct.pack('=i', 1)[0] == 0), ISBIGENDIAN)
 
@@ -783,6 +794,30 @@ class StructTest(unittest.TestCase):
         s = struct.Struct('=i2H')
         self.assertEqual(repr(s), f'Struct({s.format!r})')
 
+    @unittest.skipUnless(have_c_complex, "requires C11 complex type support")
+    def test_c_complex_round_trip(self):
+        values = [complex(*_) for _ in combinations([1, -1, 0.0, -0.0, 2,
+                                                     -3, INF, -INF, NAN], 2)]
+        for z in values:
+            for f in ['E', 'C', '>E', '>C', '<E', '<C']:
+                with self.subTest(z=z, format=f):
+                    round_trip = struct.unpack(f, struct.pack(f, z))[0]
+                    self.assertComplexesAreIdentical(z, round_trip)
+
+    @unittest.skipIf(have_c_complex, "requires no C11 complex type support")
+    def test_c_complex_error(self):
+        msg1 = "'E' format not supported on this system"
+        msg2 = "'C' format not supported on this system"
+        with self.assertRaisesRegex(struct.error, msg1):
+            struct.pack('E', 1j)
+        with self.assertRaisesRegex(struct.error, msg1):
+            struct.unpack('E', b'1')
+        with self.assertRaisesRegex(struct.error, msg2):
+            struct.pack('C', 1j)
+        with self.assertRaisesRegex(struct.error, msg2):
+            struct.unpack('C', b'1')
+
+
 class UnpackIteratorTest(unittest.TestCase):
     """
     Tests for iterative unpacking (struct.Struct.iter_unpack).