# HG changeset patch
# User Pauli Virtanen <pav@iki.fi>
# Date 1292882627 -3600
# Branch py3k
# Node ID dd49a99921679e986b171747826ec5e962c0816e
# Parent  1cc54768264b9bb752d08577826d0e7269dba920
ctypes: Fix PEP 3118 format strings for arrays

diff -r 1cc54768264b -r dd49a9992167 Lib/ctypes/test/test_pep3118.py
--- a/Lib/ctypes/test/test_pep3118.py	Mon Dec 20 20:04:51 2010 +0100
+++ b/Lib/ctypes/test/test_pep3118.py	Mon Dec 20 23:03:47 2010 +0100
@@ -92,6 +92,9 @@
 class aUnion(Union):
     _fields_ = [("a", c_int)]
 
+class StructWithArrays(Structure):
+    _fields_ = [("x", c_long * 3 * 2), ("y", Point * 4)]
+
 class Incomplete(Structure):
     pass
 
@@ -141,10 +144,10 @@
 
     ## 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",               None,           POINTER(c_short)),
 
     ## structures and unions
@@ -156,6 +159,9 @@
     (EmptyStruct,               "T{}",                  None,           EmptyStruct),
     # the pep does't support unions
     (aUnion,                    "B",                    None,           aUnion),
+    # structure with sub-arrays
+    (StructWithArrays,          "T{(2,3)<l:x:(4)T{<l:x:<l:y:}:y:}", None,  StructWithArrays),
+    (StructWithArrays * 3,      "T{(2,3)<l:x:(4)T{<l:x:<l:y:}:y:}", (3,),  StructWithArrays),
 
     ## pointer to incomplete structure
     (Incomplete,                "B",                    None,           Incomplete),
diff -r 1cc54768264b -r dd49a9992167 Modules/_ctypes/_ctypes.c
--- a/Modules/_ctypes/_ctypes.c	Mon Dec 20 20:04:51 2010 +0100
+++ b/Modules/_ctypes/_ctypes.c	Mon Dec 20 23:03:47 2010 +0100
@@ -289,6 +289,48 @@
 }
 
 /*
+  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+1) + 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, "%ld,", shape[k]);
+            } else {
+                sprintf(buf, "%ld)", 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
@@ -874,14 +916,21 @@
 
     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;
@@ -1255,7 +1304,6 @@
     long length;
     int overflow;
     Py_ssize_t itemsize, itemalign;
-    char buf[32];
 
     typedict = PyTuple_GetItem(args, 2);
     if (!typedict)
@@ -1296,13 +1344,7 @@
     }
 
     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) {
         Py_DECREF((PyObject *)stgdict);
         return NULL;
diff -r 1cc54768264b -r dd49a9992167 Modules/_ctypes/stgdict.c
--- a/Modules/_ctypes/stgdict.c	Mon Dec 20 20:04:51 2010 +0100
+++ b/Modules/_ctypes/stgdict.c	Mon Dec 20 23:03:47 2010 +0100
@@ -488,7 +488,12 @@
             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);
 
             if (stgdict->format == NULL) {