ENH: nditer: Change the Python nditer exposure to automatically add NPY_ITER_USE_MASKNA

2 files changed, 50 insertions, 5 deletions

diff --git a/numpy/core/src/multiarray/nditer_pywrap.c b/numpy/core/src/multiarray/nditer_pywrap.c
index 2e86e4753..e5061c999 100644
--- a/numpy/core/src/multiarray/nditer_pywrap.c
+++ b/numpy/core/src/multiarray/nditer_pywrap.c
@@ -613,6 +613,7 @@ npyiter_convert_ops(PyObject *op_in, PyObject *op_flags_in,
                     int *nop_out)
 {
     int iop, nop;
+    int any_maskna;
 
     /* nop and op */
     if (PyTuple_Check(op_in) || PyList_Check(op_in)) {
@@ -666,11 +667,6 @@ npyiter_convert_ops(PyObject *op_in, PyObject *op_flags_in,
             }
             else {
                 op_flags[iop] = NPY_ITER_READONLY;
-
-                /* Enable MASKNA iteration if the op needs it */
-                if (PyArray_HASMASKNA(op[iop])) {
-                    op_flags[iop] |= NPY_ITER_USE_MASKNA;
-                }
             }
         }
     }
@@ -713,6 +709,33 @@ npyiter_convert_ops(PyObject *op_in, PyObject *op_flags_in,
         }
     }
 
+    /*
+     * Because the Python exposure of nditer knows how to deal with
+     * NA-masked arrays, we automatically add NPY_ITER_USE_MASKNA
+     * flags for convenience.
+     */
+    any_maskna = 0;
+    for (iop = 0; iop < nop; ++iop) {
+        /* Enable MASKNA iteration if the op needs it */
+        if (op[iop] != NULL && PyArray_HASMASKNA(op[iop])) {
+            op_flags[iop] |= NPY_ITER_USE_MASKNA;
+            any_maskna = 1;
+        }
+    }
+    /*
+     * If any operands had an NA-mask, add it to the 'allocate' ones too.
+     * This causes the Python exposure nditer to have slightly less control
+     * than the C NpyIter usage, but is generally going to be what people
+     * want.
+     */
+    if (any_maskna) {
+        for (iop = 0; iop < nop; ++iop) {
+            if (op[iop] == NULL) {
+                op_flags[iop] |= NPY_ITER_USE_MASKNA;
+            }
+        }
+    }
+
     return 1;
 }
 
@@ -808,6 +831,7 @@ npyiter_init(NewNpyArrayIterObject *self, PyObject *args, PyObject *kwds)
         itershape.ptr = NULL;
     }
 
+
     self->iter = NpyIter_AdvancedNew(nop, op, flags, order, casting, op_flags,
                                   op_request_dtypes,
                                   oa_ndim, oa_ndim > 0 ? op_axes : NULL,
diff --git a/numpy/core/tests/test_nditer.py b/numpy/core/tests/test_nditer.py
index b13adf6d6..b0badc2d9 100644
--- a/numpy/core/tests/test_nditer.py
+++ b/numpy/core/tests/test_nditer.py
@@ -2533,5 +2533,26 @@ def test_iter_maskna():
                                  ['readonly','use_maskna'],
                                  ['readonly','arraymask']])
 
+def test_iter_maskna_default_use_maskna():
+    # The Python exposure of nditer adds the USE_MASKNA flag automatically
+    a = np.array([3, 5, np.NA, 2, 1])
+    b = np.array([1, 1.0, 4.5, 2, 0])
+
+    # The output should automatically get an NA mask
+    it = np.nditer([a,b,None])
+    for x,y,z in it:
+        z[...] = x+y
+    assert_(it.operands[2].flags.maskna)
+    assert_array_equal(it.operands[2], a+b)
+
+    # This holds even when we specify the op_flags
+    it = np.nditer([a,b.copy(),None], op_flags=[['readonly'],
+                                    ['readwrite'], ['writeonly', 'allocate']])
+    for x,y,z in it:
+        y[...] = y[...] + 1
+        z[...] = x+y
+    assert_(it.operands[2].flags.maskna)
+    assert_array_equal(it.operands[2], a+b+1)
+
 if __name__ == "__main__":
     run_module_suite()