gh-146151: Add support for complex arrays in the array module (#146237)

Co-authored-by: Victor Stinner <vstinner@python.org>

Author: skirpichev

Doc/library/array.rst

@@ -9,7 +9,7 @@
 --------------
 
 This module defines an object type which can compactly represent an array of
-basic values: characters, integers, floating-point numbers.  Arrays are mutable :term:`sequence`
+basic values: characters, integers, floating-point numbers, complex numbers.  Arrays are mutable :term:`sequence`
 types and behave very much like lists, except that the type of objects stored in
 them is constrained.  The type is specified at object creation time by using a
 :dfn:`type code`, which is a single character.  The following type codes are
@@ -46,6 +46,11 @@ defined:
 +-----------+--------------------+-------------------+-----------------------+-------+
 | ``'d'``   | double             | float             | 8                     |       |
 +-----------+--------------------+-------------------+-----------------------+-------+
+| ``'F'``   | float complex      | complex           | 8                     | \(3)  |
++-----------+--------------------+-------------------+-----------------------+-------+
+| ``'D'``   | double complex     | complex           | 16                    | \(3)  |
++-----------+--------------------+-------------------+-----------------------+-------+
+
 
 Notes:
 
@@ -63,6 +68,15 @@ Notes:
 (2)
    .. versionadded:: 3.13
 
+(3)
+   Complex types (``F`` and ``D``) are available unconditionally,
+   regardless on support for complex types (the Annex G of the C11 standard)
+   by the C compiler.
+   As specified in the C11 standard, each complex type is represented by a
+   two-element C array containing, respectively, the real and imaginary parts.
+
+   .. versionadded:: 3.15
+
 .. seealso::
 
    The :ref:`ctypes <ctypes-fundamental-data-types>` and
@@ -146,9 +160,10 @@ The module defines the following type:
    .. method:: byteswap()
 
       "Byteswap" all items of the array.  This is only supported for values which are
-      1, 2, 4, or 8 bytes in size; for other types of values, :exc:`RuntimeError` is
+      1, 2, 4, 8 or 16 bytes in size; for other types of values, :exc:`RuntimeError` is
       raised.  It is useful when reading data from a file written on a machine with a
-      different byte order.
+      different byte order.  Note, that for complex types the order of
+      components (the real part, followed by imaginary part) is preserved.
 
 
    .. method:: count(x)

Doc/whatsnew/3.15.rst

@@ -634,6 +634,14 @@ argparse
   (Contributed by Savannah Ostrowski in :gh:`142390`.)
 
 
+array
+-----
+
+* Support the :c:expr:`float complex` and :c:expr:`double complex` C types:
+  formatting characters ``'F'`` and ``'D'`` respectively.
+  (Contributed by Sergey B Kirpichev in :gh:`146151`.)
+
+
 base64
 ------
 

Lib/test/test_array.py

@@ -31,7 +31,7 @@ class ArraySubclassWithKwargs(array.array):
     def __init__(self, typecode, newarg=None):
         array.array.__init__(self)
 
-typecodes = 'uwbBhHiIlLfdqQ'
+typecodes = 'uwbBhHiIlLfdqQFD'
 
 class MiscTest(unittest.TestCase):
 
@@ -113,6 +113,12 @@ def __index__(self):
 UTF16_BE = 19
 UTF32_LE = 20
 UTF32_BE = 21
+IEEE_754_FLOAT_COMPLEX_LE = 22
+IEEE_754_FLOAT_COMPLEX_BE = 23
+IEEE_754_DOUBLE_COMPLEX_LE = 24
+IEEE_754_DOUBLE_COMPLEX_BE = 25
+
+MACHINE_FORMAT_CODE_MAX = 25
 
 
 class ArrayReconstructorTest(unittest.TestCase):
@@ -139,7 +145,7 @@ def test_error(self):
         self.assertRaises(ValueError, array_reconstructor,
                           array.array, "b", UNKNOWN_FORMAT, b"")
         self.assertRaises(ValueError, array_reconstructor,
-                          array.array, "b", 22, b"")
+                          array.array, "b", MACHINE_FORMAT_CODE_MAX + 1, b"")
         self.assertRaises(ValueError, array_reconstructor,
                           array.array, "d", 16, b"a")
 
@@ -191,7 +197,15 @@ def test_numbers(self):
             (['d'], IEEE_754_DOUBLE_LE, '<dddd',
              [9006104071832581.0, float('inf'), float('-inf'), -0.0]),
             (['d'], IEEE_754_DOUBLE_BE, '>dddd',
-             [9006104071832581.0, float('inf'), float('-inf'), -0.0])
+             [9006104071832581.0, float('inf'), float('-inf'), -0.0]),
+            (['F'], IEEE_754_FLOAT_COMPLEX_LE, '<FFFF',
+             [16711938.0j, float('inf'), complex('1-infj'), -0.0]),
+            (['F'], IEEE_754_FLOAT_COMPLEX_BE, '>FFFF',
+             [16711938.0j, float('inf'), complex('1-infj'), -0.0]),
+            (['D'], IEEE_754_DOUBLE_COMPLEX_LE, '<DDDD',
+             [9006104071832581.0j, float('inf'), complex('1-infj'), -0.0]),
+            (['D'], IEEE_754_DOUBLE_COMPLEX_BE, '>DDDD',
+             [9006104071832581.0j, float('inf'), complex('1-infj'), -0.0]),
         )
         for testcase in testcases:
             valid_typecodes, mformat_code, struct_fmt, values = testcase
@@ -279,7 +293,7 @@ def test_byteswap(self):
             example = self.example
         a = array.array(self.typecode, example)
         self.assertRaises(TypeError, a.byteswap, 42)
-        if a.itemsize in (1, 2, 4, 8):
+        if a.itemsize in (1, 2, 4, 8, 16):
             b = array.array(self.typecode, example)
             b.byteswap()
             if a.itemsize==1:
@@ -1525,6 +1539,55 @@ def test_byteswap(self):
             b.byteswap()
             self.assertEqual(a, b)
 
+class CFPTest(NumberTest):
+    example = [-42j, 0, 42+1j, 1e5j, -1e10]
+    outside = 23
+
+    def assertEntryEqual(self, entry1, entry2):
+        self.assertAlmostEqual(entry1, entry2)
+
+    def test_cmp(self):
+        a = array.array(self.typecode, self.example)
+        self.assertIs(a == 42, False)
+        self.assertIs(a != 42, True)
+
+        self.assertIs(a == a, True)
+        self.assertIs(a != a, False)
+        self.assertIs(a < a, False)
+        self.assertIs(a <= a, True)
+        self.assertIs(a > a, False)
+        self.assertIs(a >= a, True)
+
+        self.assertIs(a == 2*a, False)
+        self.assertIs(a != 2*a, True)
+        self.assertIs(a < 2*a, True)
+        self.assertIs(a <= 2*a, True)
+        self.assertIs(a > 2*a, False)
+        self.assertIs(a >= 2*a, False)
+
+    def test_nan(self):
+        a = array.array(self.typecode, [float('nan')])
+        b = array.array(self.typecode, [float('nan')])
+        self.assertIs(a != b, True)
+        self.assertIs(a == b, False)
+
+    def test_byteswap(self):
+        a = array.array(self.typecode, self.example)
+        self.assertRaises(TypeError, a.byteswap, 42)
+        if a.itemsize in (1, 2, 4, 8):
+            b = array.array(self.typecode, self.example)
+            b.byteswap()
+            if a.itemsize == 1:
+                self.assertEqual(a, b)
+            else:
+                # On alphas treating the byte swapped bit patterns as
+                # floats/doubles results in floating-point exceptions
+                # => compare the 8bit string values instead
+                self.assertNotEqual(a.tobytes(), b.tobytes())
+            b.byteswap()
+            self.assertEqual(a, b)
+
+
 class FloatTest(FPTest, unittest.TestCase):
     typecode = 'f'
     minitemsize = 4
@@ -1551,6 +1614,15 @@ def test_alloc_overflow(self):
             self.fail("Array of size > maxsize created - MemoryError expected")
 
 
+class ComplexFloatTest(CFPTest, unittest.TestCase):
+    typecode = 'F'
+    minitemsize = 8
+
+class ComplexDoubleTest(CFPTest, unittest.TestCase):
+    typecode = 'D'
+    minitemsize = 16
+
+
 class LargeArrayTest(unittest.TestCase):
     typecode = 'b'
 

Misc/NEWS.d/next/Library/2026-03-21-06-21-38.gh-issue-146151.yNpgml.rst

@@ -0,0 +1,3 @@
+Support the :c:expr:`float complex` and :c:expr:`double complex` C types
+in the :mod:`array` module: formatting characters ``'F'`` and ``'D'``
+respectively.  Patch by Sergey B Kirpichev.