ENH: dtype: Add an itemsize= keyword parameter to the dtype constructor

The 'struct'-based repr adds this parameter, as the general way to
be able to reconstruct the dtype.

3 files changed, 144 insertions, 44 deletions

diff --git a/numpy/core/src/multiarray/descriptor.c b/numpy/core/src/multiarray/descriptor.c
index 2d3b2dbbc..e071df94d 100644
--- a/numpy/core/src/multiarray/descriptor.c
+++ b/numpy/core/src/multiarray/descriptor.c
@@ -17,9 +17,6 @@
 #include "common.h"
 #include "descriptor.h"
 
-#define _chk_byteorder(arg) (arg == '>' || arg == '<' ||        \
-                             arg == '|' || arg == '=')
-
 static PyObject *typeDict = NULL;   /* Must be explicitly loaded */
 
 static PyArray_Descr *
@@ -103,6 +100,9 @@ array_set_typeDict(PyObject *NPY_UNUSED(ignored), PyObject *args)
     return Py_None;
 }
 
+#define _chk_byteorder(arg) (arg == '>' || arg == '<' ||        \
+                             arg == '|' || arg == '=')
+
 static int
 _check_for_commastring(char *type, Py_ssize_t len)
 {
@@ -136,6 +136,8 @@ _check_for_commastring(char *type, Py_ssize_t len)
     return 0;
 }
 
+#undef _chk_byteorder
+
 static int
 is_datetime_typestr(char *type, Py_ssize_t len)
 {
@@ -160,10 +162,6 @@ is_datetime_typestr(char *type, Py_ssize_t len)
     return 0;
 }
 
-
-
-#undef _chk_byteorder
-
 static PyArray_Descr *
 _convert_from_tuple(PyObject *obj)
 {
@@ -1750,22 +1748,27 @@ arraydescr_new(PyTypeObject *NPY_UNUSED(subtype),
 {
     PyObject *odescr, *ometadata=NULL;
     PyArray_Descr *descr, *conv;
-    Bool align = FALSE;
-    Bool copy = FALSE;
-    Bool copied = FALSE;
-    static char *kwlist[] = {"dtype", "align", "copy", "metadata", NULL};
+    npy_bool align = FALSE;
+    npy_bool copy = FALSE;
+    npy_bool copied = FALSE;
+    int itemsize = -1;
+
+    static char *kwlist[] = {"dtype", "align", "copy", "metadata",
+                            "itemsize", NULL};
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO&O&O!", kwlist,
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OO&O&O!i", kwlist,
                 &odescr,
                 PyArray_BoolConverter, &align,
                 PyArray_BoolConverter, &copy,
-                &PyDict_Type, &ometadata)) {
+                &PyDict_Type, &ometadata,
+                &itemsize)) {
         return NULL;
     }
 
     if ((ometadata != NULL) && (_invalid_metadata_check(ometadata))) {
         return NULL;
     }
+
     if (align) {
         if (!PyArray_DescrAlignConverter(odescr, &conv)) {
             return NULL;
@@ -1774,6 +1777,7 @@ arraydescr_new(PyTypeObject *NPY_UNUSED(subtype),
     else if (!PyArray_DescrConverter(odescr, &conv)) {
         return NULL;
     }
+
     /* Get a new copy of it unless it's already a copy */
     if (copy && conv->fields == Py_None) {
         descr = PyArray_DescrNew(conv);
@@ -1788,6 +1792,7 @@ arraydescr_new(PyTypeObject *NPY_UNUSED(subtype),
          * underlying dictionary
          */
         if (!copied) {
+            copied = TRUE;
             descr = PyArray_DescrNew(conv);
             Py_DECREF(conv);
             conv = descr;
@@ -1818,6 +1823,47 @@ arraydescr_new(PyTypeObject *NPY_UNUSED(subtype),
         }
     }
 
+    /* Override the itemsize if provided */
+    if (itemsize >= 0) {
+        /* Only can override the item size of structured arrays */
+        if (!PyDataType_HASFIELDS(conv)) {
+            PyErr_Format(PyExc_ValueError,
+                    "The itemsize parameter may only be used when "
+                    "constructing a data type with fields");
+            Py_DECREF(conv);
+            return NULL;
+        }
+        /* Make sure the itemsize isn't made too small */
+        if (itemsize < conv->elsize) {
+            PyErr_Format(PyExc_ValueError,
+                    "NumPy dtype descriptor requires %d bytes, "
+                    "cannot override to smaller itemsize of %d",
+                    (int)conv->elsize, (int)itemsize);
+            Py_DECREF(conv);
+            return NULL;
+        }
+        /* If align is set, make sure the alignment divides into the size */
+        if (align && itemsize % conv->alignment != 0) {
+            PyErr_Format(PyExc_ValueError,
+                    "NumPy dtype descriptor requires alignment of %d bytes, "
+                    "which is not divisible into the specified itemsize %d",
+                    (int)conv->alignment, (int)itemsize);
+            Py_DECREF(conv);
+            return NULL;
+        }
+        /* Change the itemsize */
+        if (itemsize != conv->elsize) {
+            /* Make sure we can change the itemsize */
+            if (!copied) {
+                copied = TRUE;
+                descr = PyArray_DescrNew(conv);
+                Py_DECREF(conv);
+                conv = descr;
+            }
+            conv->elsize = itemsize;
+        }
+    }
+
     return (PyObject *)conv;
 }
 
@@ -2771,12 +2817,20 @@ arraydescr_str(PyArray_Descr *dtype)
  * The dtype repr function specifically for structured arrays.
  */
 static PyObject *
-arraydescr_struct_repr(PyArray_Descr *self)
+arraydescr_struct_repr(PyArray_Descr *dtype)
 {
     PyObject *sub, *s;
+    int is_simple;
 
     s = PyUString_FromString("dtype(");
-    sub = arraydescr_struct_str(self);
+    if (is_dtype_struct_simple_layout(dtype)) {
+        sub = arraydescr_struct_list_str(dtype);
+        is_simple = 1;
+    }
+    else {
+        sub = arraydescr_struct_dict_str(dtype);
+        is_simple = 0;
+    }
     if (sub == NULL) {
         return NULL;
     }
@@ -2784,9 +2838,14 @@ arraydescr_struct_repr(PyArray_Descr *self)
     PyUString_ConcatAndDel(&s, sub);
 
     /* If it's an aligned structure, add the align=True parameter */
-    if (self->flags&NPY_ALIGNED_STRUCT) {
+    if (dtype->flags&NPY_ALIGNED_STRUCT) {
         PyUString_ConcatAndDel(&s, PyUString_FromString(", align=True"));
     }
+    /* If it wasn't simple, also specify the itemsize parameter */
+    if (!is_simple) {
+        PyUString_ConcatAndDel(&s,
+                    PyUString_FromFormat(", itemsize=%d", (int)dtype->elsize));
+    }
 
     PyUString_ConcatAndDel(&s, PyUString_FromString(")"));
     return s;
diff --git a/numpy/core/tests/test_dtype.py b/numpy/core/tests/test_dtype.py
index c79b755be..83df4b7ed 100644
--- a/numpy/core/tests/test_dtype.py
+++ b/numpy/core/tests/test_dtype.py
@@ -41,6 +41,16 @@ class TestBuiltin(TestCase):
             #print typestr
             assert_raises(TypeError, np.dtype, typestr)
 
+    def test_bad_param(self):
+        # Can't override the itemsize of a non-struct type
+        assert_raises(ValueError, np.dtype, 'f4', itemsize=6)
+        # Even if you specify the same size
+        assert_raises(ValueError, np.dtype, 'f4', itemsize=4)
+        # Can't give a size that's too small
+        assert_raises(ValueError, np.dtype, 'f4, i4', itemsize=7)
+        # If alignment is enabled, the alignment (4) must divide the itemsize
+        assert_raises(ValueError, np.dtype, 'f4, i1', align=True, itemsize=9)
+
 class TestRecord(TestCase):
     def test_equivalent_record(self):
         """Test whether equivalent record dtypes hash the same."""
@@ -165,5 +175,64 @@ class TestMetadata(TestCase):
         d = np.dtype([('a', np.dtype(int, metadata={'datum': 1}))])
         self.assertEqual(d['a'].metadata, {'datum': 1})
 
+class TestString(TestCase):
+    def test_complex_dtype_str(self):
+        dt = np.dtype([('top', [('tiles', ('>f4', (64, 64)), (1,)),
+                                ('rtile', '>f4', (64, 36))], (3,)),
+                       ('bottom', [('bleft', ('>f4', (8, 64)), (1,)),
+                                   ('bright', '>f4', (8, 36))])])
+        assert_equal(str(dt),
+                     "[('top', [('tiles', ('>f4', (64, 64)), (1,)), "
+                     "('rtile', '>f4', (64, 36))], (3,)), "
+                     "('bottom', [('bleft', ('>f4', (8, 64)), (1,)), "
+                     "('bright', '>f4', (8, 36))])]")
+
+        dt = np.dtype({'names': ['r','g','b'], 'formats': ['u1', 'u1', 'u1'],
+                        'offsets': [0, 1, 2],
+                        'titles': ['Red pixel', 'Green pixel', 'Blue pixel']})
+        assert_equal(str(dt),
+                    "[(('Red pixel', 'r'), 'u1'), "
+                    "(('Green pixel', 'g'), 'u1'), "
+                    "(('Blue pixel', 'b'), 'u1')]")
+
+        dt = np.dtype({'names': ['r','b'], 'formats': ['u1', 'u1'],
+                        'offsets': [0, 2],
+                        'titles': ['Red pixel', 'Blue pixel']})
+        assert_equal(str(dt),
+                    "{'names':['r','b'], "
+                    "'formats':['u1','u1'], "
+                    "'offsets':[0,2], "
+                    "'titles':['Red pixel','Blue pixel']}")
+
+    def test_complex_dtype_repr(self):
+        dt = np.dtype([('top', [('tiles', ('>f4', (64, 64)), (1,)),
+                                ('rtile', '>f4', (64, 36))], (3,)),
+                       ('bottom', [('bleft', ('>f4', (8, 64)), (1,)),
+                                   ('bright', '>f4', (8, 36))])])
+        assert_equal(repr(dt),
+                     "dtype([('top', [('tiles', ('>f4', (64, 64)), (1,)), "
+                     "('rtile', '>f4', (64, 36))], (3,)), "
+                     "('bottom', [('bleft', ('>f4', (8, 64)), (1,)), "
+                     "('bright', '>f4', (8, 36))])])")
+
+        dt = np.dtype({'names': ['r','g','b'], 'formats': ['u1', 'u1', 'u1'],
+                        'offsets': [0, 1, 2],
+                        'titles': ['Red pixel', 'Green pixel', 'Blue pixel']},
+                        align=True)
+        assert_equal(repr(dt),
+                    "dtype([(('Red pixel', 'r'), 'u1'), "
+                    "(('Green pixel', 'g'), 'u1'), "
+                    "(('Blue pixel', 'b'), 'u1')], align=True)")
+
+        dt = np.dtype({'names': ['r','b'], 'formats': ['u1', 'u1'],
+                        'offsets': [0, 2],
+                        'titles': ['Red pixel', 'Blue pixel']},
+                        itemsize = 4)
+        assert_equal(repr(dt),
+                    "dtype({'names':['r','b'], "
+                    "'formats':['u1','u1'], "
+                    "'offsets':[0,2], "
+                    "'titles':['Red pixel','Blue pixel']}, itemsize=4)")
+
 if __name__ == "__main__":
     run_module_suite()
diff --git a/numpy/core/tests/test_regression.py b/numpy/core/tests/test_regression.py
index a51f3cee4..1c43189bb 100644
--- a/numpy/core/tests/test_regression.py
+++ b/numpy/core/tests/test_regression.py
@@ -853,34 +853,6 @@ class TestRegression(TestCase):
         assert_almost_equal(fdouble, 1.234)
         assert_almost_equal(flongdouble, 1.234)
 
-    def test_complex_dtype_printing(self, level=rlevel):
-        dt = np.dtype([('top', [('tiles', ('>f4', (64, 64)), (1,)),
-                                ('rtile', '>f4', (64, 36))], (3,)),
-                       ('bottom', [('bleft', ('>f4', (8, 64)), (1,)),
-                                   ('bright', '>f4', (8, 36))])])
-        assert_equal(str(dt),
-                     "[('top', [('tiles', ('>f4', (64, 64)), (1,)), "
-                     "('rtile', '>f4', (64, 36))], (3,)), "
-                     "('bottom', [('bleft', ('>f4', (8, 64)), (1,)), "
-                     "('bright', '>f4', (8, 36))])]")
-
-        dt = np.dtype({'names': ['r','g','b'], 'formats': ['u1', 'u1', 'u1'],
-                        'offsets': [0, 1, 2],
-                        'titles': ['Red pixel', 'Green pixel', 'Blue pixel']})
-        assert_equal(str(dt),
-                    "[(('Red pixel', 'r'), 'u1'), "
-                    "(('Green pixel', 'g'), 'u1'), "
-                    "(('Blue pixel', 'b'), 'u1')]")
-
-        dt = np.dtype({'names': ['r','b'], 'formats': ['u1', 'u1'],
-                        'offsets': [0, 2],
-                        'titles': ['Red pixel', 'Blue pixel']})
-        assert_equal(str(dt),
-                    "{'names':['r','b'], "
-                    "'formats':['u1','u1'], "
-                    "'offsets':[0,2], "
-                    "'titles':['Red pixel','Blue pixel']}")
-
     def test_nonnative_endian_fill(self, level=rlevel):
         """ Non-native endian arrays were incorrectly filled with scalars before
         r5034.