BUG: nditer: NA masks in arrays with leading 1 dimensions had an issue

6 files changed, 25 insertions, 9 deletions

diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py
index f2470a58f..042688fc2 100644
--- a/numpy/core/code_generators/numpy_api.py
+++ b/numpy/core/code_generators/numpy_api.py
@@ -335,6 +335,7 @@ multiarray_funcs_api = {
     'PyArray_CreateReduceResult':           296,
     'PyArray_InitializeReduceResult':       297,
     'PyArray_RemoveAxesInPlace':            298,
+    'PyArray_DebugPrint':                   299,
 }
 
 ufunc_types_api = {
diff --git a/numpy/core/fromnumeric.py b/numpy/core/fromnumeric.py
index 4d66b51b2..12cdbd68a 100644
--- a/numpy/core/fromnumeric.py
+++ b/numpy/core/fromnumeric.py
@@ -877,6 +877,8 @@ def squeeze(a, axis=None):
     a : array_like
         Input data.
     axis : None or int or tuple of ints, optional
+        .. versionadded:: 1.7.0
+
         Selects a subset of the single-dimensional entries in the
         shape. If an axis is selected with shape entry greater than
         one, an error is raised.
diff --git a/numpy/core/src/multiarray/arrayobject.c b/numpy/core/src/multiarray/arrayobject.c
index 17cbd328d..d21a78d74 100644
--- a/numpy/core/src/multiarray/arrayobject.c
+++ b/numpy/core/src/multiarray/arrayobject.c
@@ -454,6 +454,10 @@ dump_data(char **string, int *n, int *max_n, char *data, int nd,
 #undef CHECK_MEMORY
 }
 
+/*NUMPY_API
+ * Prints the raw data of the ndarray in a form useful for debugging
+ * low-level C issues.
+ */
 NPY_NO_EXPORT void
 PyArray_DebugPrint(PyArrayObject *obj)
 {
diff --git a/numpy/core/src/multiarray/arrayobject.h b/numpy/core/src/multiarray/arrayobject.h
index b4c4c086c..ec3361435 100644
--- a/numpy/core/src/multiarray/arrayobject.h
+++ b/numpy/core/src/multiarray/arrayobject.h
@@ -12,7 +12,4 @@ _strings_richcompare(PyArrayObject *self, PyArrayObject *other, int cmp_op,
 NPY_NO_EXPORT PyObject *
 array_richcompare(PyArrayObject *self, PyObject *other, int cmp_op);
 
-NPY_NO_EXPORT void
-PyArray_DebugPrint(PyArrayObject *obj);
-
 #endif
diff --git a/numpy/core/src/multiarray/nditer_constr.c b/numpy/core/src/multiarray/nditer_constr.c
index 93c11b60d..e16738f7e 100644
--- a/numpy/core/src/multiarray/nditer_constr.c
+++ b/numpy/core/src/multiarray/nditer_constr.c
@@ -3092,7 +3092,7 @@ npyiter_allocate_arrays(NpyIter *iter,
         }
 
         /*
-         * If no alignment, byte swap, or casting is needed, and
+         * If no alignment, byte swap, or casting is needed,
          * the inner stride of this operand works for the whole
          * array, we can set NPY_OP_ITFLAG_BUFNEVER.
          */
@@ -3301,12 +3301,19 @@ npyiter_fill_maskna_axisdata(NpyIter *iter, int **op_axes)
     /* Initialize the strides of any BUFNEVER mask operands */
     if (itflags & NPY_ITFLAG_BUFFER) {
         npy_intp *strides = NBF_STRIDES(bufferdata);
-        axisdata = NIT_AXISDATA(iter);
 
         for (iop = first_maskna_op; iop < nop; ++iop) {
             if (op_itflags[iop] & NPY_OP_ITFLAG_BUFNEVER) {
                 if (PyArray_HASMASKNA(op[iop])) {
-                    strides[iop] = NAD_STRIDES(axisdata)[iop];
+                    axisdata = NIT_AXISDATA(iter);
+                    /* Find stride of the first non-empty shape */
+                    for (idim = 0; idim < ndim; ++idim) {
+                        if (NAD_SHAPE(axisdata) != 1) {
+                            strides[iop] = NAD_STRIDES(axisdata)[iop];
+                            break;
+                        }
+                        NIT_ADVANCE_AXISDATA(axisdata, 1);
+                    }
                 }
                 else {
                     strides[iop] = 0;
diff --git a/numpy/core/tests/test_maskna.py b/numpy/core/tests/test_maskna.py
index 8454db4f9..ad80f403d 100644
--- a/numpy/core/tests/test_maskna.py
+++ b/numpy/core/tests/test_maskna.py
@@ -694,9 +694,14 @@ def test_maskna_ufunc_1D():
     assert_raises(ValueError, np.add, a, b, out=c_orig)
 
     # Divide in-place with NA
-    a = np.array([[np.NA], [12.]], maskna=True)
-    assert_array_equal(a, [[np.NA], [666.]])
-    # The assert_array_equal should have caught this...
+    a_orig = np.array([[3], [12.]])
+    a = a_orig.view(maskna=True)
+    a[0,0] = np.NA
+    a /= 3
+    # Shouldn't have touched the masked element
+    assert_array_equal(a_orig, [[3], [4.]])
+    assert_array_equal(a, [[np.NA], [4.]])
+    # double-check assertions
     assert_equal(np.isna(a), [[1], [0]])
     assert_equal(a[~np.isna(a)], [4.])