Merge pull request #3725 from juliantaylor/fix-subarray-zeros

BUG: fix zero initialization with subarray dtypes

2 files changed, 42 insertions, 27 deletions

diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
index ef0dca414..78616116d 100644
--- a/numpy/core/src/multiarray/ctors.c
+++ b/numpy/core/src/multiarray/ctors.c
@@ -863,15 +863,16 @@ discover_dimensions(PyObject *obj, int *maxndim, npy_intp *d, int check_it,
     return 0;
 }
 
-/*NUMPY_API
+/*
  * Generic new array creation routine.
+ * Internal variant with calloc argument for PyArray_Zeros.
  *
  * steals a reference to descr (even on failure)
  */
-NPY_NO_EXPORT PyObject *
-PyArray_NewFromDescr(PyTypeObject *subtype, PyArray_Descr *descr, int nd,
-                     npy_intp *dims, npy_intp *strides, void *data,
-                     int flags, PyObject *obj)
+static PyObject *
+PyArray_NewFromDescr_int(PyTypeObject *subtype, PyArray_Descr *descr, int nd,
+                         npy_intp *dims, npy_intp *strides, void *data,
+                         int flags, PyObject *obj, int zeroed)
 {
     PyArrayObject_fields *fa;
     int i;
@@ -890,9 +891,9 @@ PyArray_NewFromDescr(PyTypeObject *subtype, PyArray_Descr *descr, int nd,
         }
         nd =_update_descr_and_dimensions(&descr, newdims,
                                          newstrides, nd);
-        ret = PyArray_NewFromDescr(subtype, descr, nd, newdims,
-                                   newstrides,
-                                   data, flags, obj);
+        ret = PyArray_NewFromDescr_int(subtype, descr, nd, newdims,
+                                       newstrides,
+                                       data, flags, obj, zeroed);
         return ret;
     }
 
@@ -1026,7 +1027,7 @@ PyArray_NewFromDescr(PyTypeObject *subtype, PyArray_Descr *descr, int nd,
          * It is bad to have unitialized OBJECT pointers
          * which could also be sub-fields of a VOID array
          */
-        if (PyDataType_FLAGCHK(descr, NPY_NEEDS_INIT)) {
+        if (zeroed || PyDataType_FLAGCHK(descr, NPY_NEEDS_INIT)) {
             data = PyDataMem_NEW_ZEROED(sd, 1);
         }
         else {
@@ -1102,6 +1103,22 @@ PyArray_NewFromDescr(PyTypeObject *subtype, PyArray_Descr *descr, int nd,
     return NULL;
 }
 
+
+/*NUMPY_API
+ * Generic new array creation routine.
+ *
+ * steals a reference to descr (even on failure)
+ */
+NPY_NO_EXPORT PyObject *
+PyArray_NewFromDescr(PyTypeObject *subtype, PyArray_Descr *descr, int nd,
+                     npy_intp *dims, npy_intp *strides, void *data,
+                     int flags, PyObject *obj)
+{
+    return PyArray_NewFromDescr_int(subtype, descr, nd,
+                                    dims, strides, data,
+                                    flags, obj, 0);
+}
+
 /*NUMPY_API
  * Creates a new array with the same shape as the provided one,
  * with possible memory layout order and data type changes.
@@ -2722,22 +2739,16 @@ NPY_NO_EXPORT PyObject *
 PyArray_Zeros(int nd, npy_intp *dims, PyArray_Descr *type, int is_f_order)
 {
     PyArrayObject *ret;
-    int need_init;
 
     if (!type) {
         type = PyArray_DescrFromType(NPY_DEFAULT_TYPE);
     }
 
-    /* set init flag on copy of the descriptor to get zero memory */
-    PyArray_DESCR_REPLACE(type);
-    need_init = PyDataType_FLAGCHK(type, NPY_NEEDS_INIT);
-    type->flags |= NPY_NEEDS_INIT;
-
-    ret = (PyArrayObject *)PyArray_NewFromDescr(&PyArray_Type,
-                                                type,
-                                                nd, dims,
-                                                NULL, NULL,
-                                                is_f_order, NULL);
+    ret = (PyArrayObject *)PyArray_NewFromDescr_int(&PyArray_Type,
+                                                    type,
+                                                    nd, dims,
+                                                    NULL, NULL,
+                                                    is_f_order, NULL, 1);
 
     if (ret == NULL) {
         return NULL;
@@ -2751,13 +2762,6 @@ PyArray_Zeros(int nd, npy_intp *dims, PyArray_Descr *type, int is_f_order)
         }
     }
 
-    /*
-     * drop init flag again if its not needed,
-     * e.g. avoids unnecessary initializations of nditer buffers
-     */
-    if (!need_init) {
-        PyArray_DESCR(ret)->flags &= ~NPY_NEEDS_INIT;
-    }
 
     return (PyObject *)ret;
 
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index 19d2810a2..77da5543d 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -448,6 +448,17 @@ class TestCreation(TestCase):
             assert_equal(np.count_nonzero(d), 0)
             assert_equal(d.sum(), 0)
 
+            d = np.zeros(2, dtype='(2,4)i4')
+            assert_equal(np.count_nonzero(d), 0)
+            assert_equal(d.sum(), 0)
+
+            d = np.zeros(2, dtype='4i4')
+            assert_equal(np.count_nonzero(d), 0)
+            assert_equal(d.sum(), 0)
+
+            d = np.zeros(2, dtype='(2,4)i4, (2,4)i4')
+            assert_equal(np.count_nonzero(d), 0)
+
     def test_zeros_obj(self):
         # test initialization from PyLong(0)
         d = np.zeros((13,), dtype=object)