support pep 3118 format strings for ctypes objects with nontrivial shapes (closes #10744)

Patch by Matti Picus.

5 files changed, 74 insertions, 15 deletions

diff --git a/Lib/ctypes/test/test_pep3118.py b/Lib/ctypes/test/test_pep3118.py
index ad13b01..32f802c 100644
--- a/Lib/ctypes/test/test_pep3118.py
+++ b/Lib/ctypes/test/test_pep3118.py
@@ -96,6 +96,9 @@ class EmptyStruct(Structure):
 class aUnion(Union):
     _fields_ = [("a", c_int)]
 
+class StructWithArrays(Structure):
+    _fields_ = [("x", c_long * 3 * 2), ("y", Point * 4)]
+
 class Incomplete(Structure):
     pass
 
@@ -145,10 +148,10 @@ native_types = [
 
     ## arrays and pointers
 
-    (c_double * 4,              "(4)<d",                (4,),           c_double),
-    (c_float * 4 * 3 * 2,       "(2,3,4)<f",            (2,3,4),        c_float),
-    (POINTER(c_short) * 2,      "(2)&<h",               (2,),           POINTER(c_short)),
-    (POINTER(c_short) * 2 * 3,  "(3,2)&<h",             (3,2,),         POINTER(c_short)),
+    (c_double * 4,              "<d",                   (4,),           c_double),
+    (c_float * 4 * 3 * 2,       "<f",                   (2,3,4),        c_float),
+    (POINTER(c_short) * 2,      "&<h",                  (2,),           POINTER(c_short)),
+    (POINTER(c_short) * 2 * 3,  "&<h",                  (3,2,),         POINTER(c_short)),
     (POINTER(c_short * 2),      "&(2)<h",               (),           POINTER(c_short)),
 
     ## structures and unions
@@ -160,6 +163,9 @@ native_types = [
     (EmptyStruct,               "T{}",                  (),           EmptyStruct),
     # the pep does't support unions
     (aUnion,                    "B",                    (),           aUnion),
+    # structure with sub-arrays
+    (StructWithArrays,          "T{(2,3)<l:x:(4)T{<l:x:<l:y:}:y:}", (),  StructWithArrays),
+    (StructWithArrays * 3,      "T{(2,3)<l:x:(4)T{<l:x:<l:y:}:y:}", (3,),  StructWithArrays),
 
     ## pointer to incomplete structure
     (Incomplete,                "B",                    (),           Incomplete),
diff --git a/Misc/NEWS b/Misc/NEWS
index e2836a1..9e2de7a 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -23,6 +23,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #10744: Fix PEP 3118 format strings on ctypes objects with a nontrivial
+  shape.
+
 - Issue #20998: Fixed re.fullmatch() of repeated single character pattern
   with ignore case.  Original patch by Matthew Barnett.
 
diff --git a/Modules/_ctypes/_ctypes.c b/Modules/_ctypes/_ctypes.c
index 97ae798..01adbb5 100644
--- a/Modules/_ctypes/_ctypes.c
+++ b/Modules/_ctypes/_ctypes.c
@@ -289,6 +289,48 @@ _ctypes_alloc_format_string(const char *prefix, const char *suffix)
 }
 
 /*
+  Allocate a memory block for a pep3118 format string, adding
+  the given prefix (if non-null), an additional shape prefix, and a suffix.
+  Returns NULL on failure, with the error indicator set.  If called with
+  a suffix of NULL the error indicator must already be set.
+ */
+char *
+_ctypes_alloc_format_string_with_shape(int ndim, const Py_ssize_t *shape,
+                                       const char *prefix, const char *suffix)
+{
+    char *new_prefix;
+    char *result;
+    char buf[32];
+    int prefix_len;
+    int k;
+
+    prefix_len = 32 * ndim + 3;
+    if (prefix)
+        prefix_len += strlen(prefix);
+    new_prefix = PyMem_Malloc(prefix_len);
+    if (new_prefix == NULL)
+        return NULL;
+    new_prefix[0] = '\0';
+    if (prefix)
+        strcpy(new_prefix, prefix);
+    if (ndim > 0) {
+        /* Add the prefix "(shape[0],shape[1],...,shape[ndim-1])" */
+        strcat(new_prefix, "(");
+        for (k = 0; k < ndim; ++k) {
+            if (k < ndim-1) {
+                sprintf(buf, "%"PY_FORMAT_SIZE_T"d,", shape[k]);
+            } else {
+                sprintf(buf, "%"PY_FORMAT_SIZE_T"d)", shape[k]);
+            }
+            strcat(new_prefix, buf);
+        }
+    }
+    result = _ctypes_alloc_format_string(new_prefix, suffix);
+    PyMem_Free(new_prefix);
+    return result;
+}
+
+/*
   PyCStructType_Type - a meta type/class.  Creating a new class using this one as
   __metaclass__ will call the contructor StructUnionType_new.  It replaces the
   tp_dict member with a new instance of StgDict, and initializes the C
@@ -860,14 +902,21 @@ PyCPointerType_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 
     if (proto) {
         StgDictObject *itemdict = PyType_stgdict(proto);
+        const char *current_format;
         assert(itemdict);
         /* If itemdict->format is NULL, then this is a pointer to an
            incomplete type.  We create a generic format string
            'pointer to bytes' in this case.  XXX Better would be to
            fix the format string later...
         */
-        stgdict->format = _ctypes_alloc_format_string("&",
-                      itemdict->format ? itemdict->format : "B");
+        current_format = itemdict->format ? itemdict->format : "B";
+        if (itemdict->shape != NULL) {
+            /* pointer to an array: the shape needs to be prefixed */
+            stgdict->format = _ctypes_alloc_format_string_with_shape(
+                itemdict->ndim, itemdict->shape, "&", current_format);
+        } else {
+            stgdict->format = _ctypes_alloc_format_string("&", current_format);
+        }
         if (stgdict->format == NULL) {
             Py_DECREF((PyObject *)stgdict);
             return NULL;
@@ -1245,7 +1294,6 @@ PyCArrayType_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
     long length;
     int overflow;
     Py_ssize_t itemsize, itemalign;
-    char buf[32];
 
     /* create the new instance (which is a class,
        since we are a metatype!) */
@@ -1295,13 +1343,7 @@ PyCArrayType_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
     }
 
     assert(itemdict->format);
-    if (itemdict->format[0] == '(') {
-        sprintf(buf, "(%ld,", length);
-        stgdict->format = _ctypes_alloc_format_string(buf, itemdict->format+1);
-    } else {
-        sprintf(buf, "(%ld)", length);
-        stgdict->format = _ctypes_alloc_format_string(buf, itemdict->format);
-    }
+    stgdict->format = _ctypes_alloc_format_string(NULL, itemdict->format);
     if (stgdict->format == NULL)
         goto error;
     stgdict->ndim = itemdict->ndim + 1;
diff --git a/Modules/_ctypes/ctypes.h b/Modules/_ctypes/ctypes.h
index 5237ac2..ac8341e 100644
--- a/Modules/_ctypes/ctypes.h
+++ b/Modules/_ctypes/ctypes.h
@@ -357,6 +357,9 @@ extern void _ctypes_add_traceback(char *, char *, int);
 
 extern PyObject *PyCData_FromBaseObj(PyObject *type, PyObject *base, Py_ssize_t index, char *adr);
 extern char *_ctypes_alloc_format_string(const char *prefix, const char *suffix);
+extern char *_ctypes_alloc_format_string_with_shape(int ndim,
+                                                const Py_ssize_t *shape,
+                                                const char *prefix, const char *suffix);
 
 extern int _ctypes_simple_instance(PyObject *obj);
 
diff --git a/Modules/_ctypes/stgdict.c b/Modules/_ctypes/stgdict.c
index b95b0a4..728f751 100644
--- a/Modules/_ctypes/stgdict.c
+++ b/Modules/_ctypes/stgdict.c
@@ -505,7 +505,12 @@ PyCStructUnionType_update_stgdict(PyObject *type, PyObject *fields, int isStruct
             sprintf(buf, "%s:%s:", fieldfmt, fieldname);
 
             ptr = stgdict->format;
-            stgdict->format = _ctypes_alloc_format_string(stgdict->format, buf);
+            if (dict->shape != NULL) {
+                stgdict->format = _ctypes_alloc_format_string_with_shape(
+                    dict->ndim, dict->shape, stgdict->format, buf);
+            } else {
+                stgdict->format = _ctypes_alloc_format_string(stgdict->format, buf);
+            }
             PyMem_Free(ptr);
             PyMem_Free(buf);