test_buffer: test also Zf and Zd formats

Lib/test/test_buffer.py

@@ -67,7 +67,7 @@
     'h':0, 'H':0, 'i':0, 'I':0,
     'l':0, 'L':0, 'n':0, 'N':0,
     'e':0, 'f':0, 'd':0, 'P':0,
-    'F':0, 'D':0
+    'F':0, 'D':0, 'Zf':0, 'Zd':0,
 }
 
 # NumPy does not have 'n' or 'N':
@@ -95,7 +95,9 @@
     'e':(-65519, 65520),   'f':(-(1<<63), 1<<63),
     'd':(-(1<<1023), 1<<1023),
     'F':(-(1<<63), 1<<63),
-    'D':(-(1<<1023), 1<<1023)
+    'D':(-(1<<1023), 1<<1023),
+    'Zf':(-(1<<63), 1<<63),
+    'Zd':(-(1<<1023), 1<<1023),
 }
 
 def native_type_range(fmt):
@@ -110,9 +112,9 @@ def native_type_range(fmt):
         lh = (-(1<<63), 1<<63)
     elif fmt == 'd':
         lh = (-(1<<1023), 1<<1023)
-    elif fmt == 'F':
+    elif fmt in ('F', 'Zf'):
         lh = (-(1<<63), 1<<63)
-    elif fmt == 'D':
+    elif fmt in ('D', 'Zd'):
         lh = (-(1<<1023), 1<<1023)
     else:
         for exp in (128, 127, 64, 63, 32, 31, 16, 15, 8, 7):
@@ -182,18 +184,28 @@ def randrange_fmt(mode, char, obj):
     if char in 'efd':
         x = struct.pack(char, x)
         x = struct.unpack(char, x)[0]
-    if char in 'FD':
+    if char in ('F', 'D', 'Zf', 'Zd'):
         y = randrange(*fmtdict[mode][char])
         x = complex(x, y)
         x = struct.pack(char, x)
         x = struct.unpack(char, x)[0]
     return x
 
+def split_format(fmt):
+    i = 0
+    while i < len(fmt):
+        if fmt[i] == 'Z':
+            n = 2
+        else:
+            n = 1
+        yield fmt[i:i + n]
+        i += n
+
 def gen_item(fmt, obj):
     """Return single random item."""
     mode, chars = fmt.split('#')
     x = []
-    for c in chars:
+    for c in split_format(chars):
         x.append(randrange_fmt(mode, c, obj))
     return x[0] if len(x) == 1 else tuple(x)
 
@@ -254,9 +266,7 @@ def is_byte_format(fmt):
 
 def is_memoryview_format(fmt):
     """format suitable for memoryview"""
-    x = len(fmt)
-    return ((x == 1 or (x == 2 and fmt[0] == '@')) and
-            fmt[x-1] in MEMORYVIEW)
+    return fmt.removeprefix('@') in MEMORYVIEW
 
 NON_BYTE_FORMAT = [c for c in fmtdict['@'] if not is_byte_format(c)]
 
@@ -648,14 +658,22 @@ def ndarray_from_structure(items, fmt, t, flags=0):
     return ndarray(items, shape=shape, strides=strides, format=fmt,
                    offset=offset, flags=ND_WRITABLE|flags)
 
+# Convert PEP 3118 formats to numpy dtypes
+FORMAT_TO_DTYPE = {
+    'Zf': 'F',
+    'Zd': 'D',
+}
+
 def numpy_array_from_structure(items, fmt, t):
     """Return numpy_array from the tuple returned by rand_structure()"""
     memlen, itemsize, ndim, shape, strides, offset = t
     buf = bytearray(memlen)
     for j, v in enumerate(items):
         struct.pack_into(fmt, buf, j*itemsize, v)
+    # Replace Zd/Zf formats with D/F dtypes
+    dtype = FORMAT_TO_DTYPE.get(fmt, fmt)
     return numpy_array(buffer=buf, shape=shape, strides=strides,
-                       dtype=fmt, offset=offset)
+                       dtype=dtype, offset=offset)
 
 
 # ======================================================================
@@ -3037,7 +3055,7 @@ def test_memoryview_assign(self):
                 continue
             m2 = m1.cast(fmt)
             lo, hi = _range
-            if fmt in "dfDF":
+            if fmt in ("d", "f", "D", "F", "Zd", "Zf"):
                 lo, hi = -2**1024, 2**1024
             if fmt != 'P': # PyLong_AsVoidPtr() accepts negative numbers
                 self.assertRaises(ValueError, m2.__setitem__, 0, lo-1)