ENH: missingdata: More progress towards NPY_ITER_USE_MASKNA flag support

3 files changed, 61 insertions, 8 deletions

diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py
index a28ecdd19..3dd32cf01 100644
--- a/numpy/core/numeric.py
+++ b/numpy/core/numeric.py
@@ -1330,15 +1330,15 @@ def array_repr(arr, max_line_width=None, precision=None, suppress_small=None):
     else:
         cName = "array"
 
-    skiptype = (arr.dtype.type in _typelessdata) and arr.size > 0
+    skipdtype = (arr.dtype.type in _typelessdata) and arr.size > 0
 
     if arr.flags.maskna:
         lst += ", maskna=True"
         # If everything is NA, can't skip the type
-        if skiptype and np.all(np.isna(arr)):
-            skiptype = False
+        if skipdtype and all(isna(arr)):
+            skipdtype = False
 
-    if skiptype:
+    if skipdtype:
         return "%s(%s)" % (cName, lst)
     else:
         typename = arr.dtype.name
diff --git a/numpy/core/src/multiarray/na_mask.c b/numpy/core/src/multiarray/na_mask.c
index eb3256cdf..f374420ac 100644
--- a/numpy/core/src/multiarray/na_mask.c
+++ b/numpy/core/src/multiarray/na_mask.c
@@ -273,8 +273,7 @@ PyArray_IsNA(PyObject *obj)
             return NULL;
         }
 
-        PyErr_SetString(PyExc_RuntimeError,
-                "isna isn't done yet");
+        PyErr_SetString(PyExc_RuntimeError, "isna isn't done yet");
         return NULL;
 
         dtype = PyArray_DescrFromType(NPY_BOOL);
diff --git a/numpy/core/src/multiarray/nditer_constr.c b/numpy/core/src/multiarray/nditer_constr.c
index 37842c72d..2054a4fdb 100644
--- a/numpy/core/src/multiarray/nditer_constr.c
+++ b/numpy/core/src/multiarray/nditer_constr.c
@@ -1029,11 +1029,12 @@ npyiter_prepare_one_operand(PyArrayObject **op,
         return 0;
     }
 
+
     if (PyArray_Check(*op)) {
         if (((*op_itflags) & NPY_OP_ITFLAG_WRITE) &&
                     (!PyArray_CHKFLAGS(*op, NPY_ARRAY_WRITEABLE))) {
             PyErr_SetString(PyExc_ValueError,
-                    "Iterator operand was a non-writeable array, but was "
+                    "Operand was a non-writeable array, but "
                     "flagged as writeable");
             return 0;
         }
@@ -1042,6 +1043,25 @@ npyiter_prepare_one_operand(PyArrayObject **op,
                     "Iteration of zero-sized operands is not enabled");
             return 0;
         }
+        /*
+         * Writeable USE_MASKNA operands must have a mask
+         * (or NA dtype, later)
+         */
+        if ((op_flags & NPY_ITER_USE_MASKNA) != 0 &&
+                            ((*op_itflags) & NPY_OP_ITFLAG_WRITE) != 0 &&
+                            !PyArray_HASMASKNA(*op)) {
+            PyErr_SetString(PyExc_ValueError,
+                    "Operand is writeable and flagged USE_MASKNA, "
+                    "but the operand does not have an NA mask");
+            return 0;
+        }
+        /* Arrays with NA masks must have USE_MASKNA specified */
+        if ((op_flags & NPY_ITER_USE_MASKNA) == 0 && PyArray_HASMASKNA(*op)) {
+            PyErr_SetString(PyExc_ValueError,
+                    "Operand has an NA mask but the operation does "
+                    "not support NA via the flag USE_MASKNA");
+            return 0;
+        }
         *op_dataptr = PyArray_BYTES(*op);
         /* PyArray_DESCR does not give us a reference */
         *op_dtype = PyArray_DESCR(*op);
@@ -1236,6 +1256,13 @@ npyiter_prepare_operands(int nop, int first_maskna_op, PyArrayObject **op_in,
         }
     }
 
+    /* Initialize the mask virtual operands to NULL for now */
+    for (iop = first_maskna_op; iop < nop; ++iop) {
+        op[iop] = NULL;
+        op_dataptr[iop] = NULL;
+        op_dtype[iop] = NULL;
+    }
+
     /* If all the operands were NULL, it's an error */
     if (op[0] == NULL) {
         int all_null = 1;
@@ -2842,6 +2869,17 @@ npyiter_allocate_arrays(NpyIter *iter,
             op[iop] = out;
 
             /*
+             * Add an NA mask to the array if needed. When NA dtypes
+             * are supported, this should skip allocating the mask
+             * if the allocated array has an NA dtype.
+             */
+            if (op_flags[iop] & NPY_ITER_USE_MASKNA) {
+                if (PyArray_AllocateMaskNA(out, 1, 0) < 0) {
+                    return 0;
+                }
+            }
+
+            /*
              * Now we need to replace the pointers and strides with values
              * from the new array.
              */
@@ -2870,6 +2908,12 @@ npyiter_allocate_arrays(NpyIter *iter,
             if (temp == NULL) {
                 return 0;
             }
+            /* Add an NA mask if needed */
+            if (PyArray_HASMASKNA(op[iop])) {
+                if (PyArray_AllocateMaskNA(temp, 1, 0) < 0) {
+                    return 0;
+                }
+            }
             if (PyArray_CopyInto(temp, op[iop]) != 0) {
                 Py_DECREF(temp);
                 return 0;
@@ -2911,8 +2955,18 @@ npyiter_allocate_arrays(NpyIter *iter,
             if (temp == NULL) {
                 return 0;
             }
+            /* Add an NA mask if needed */
+            if (PyArray_HASMASKNA(op[iop])) {
+                if (PyArray_AllocateMaskNA(temp, 1, 0) < 0) {
+                    return 0;
+                }
+            }
 
-            /* If the data will be read, copy it into temp */
+            /*
+             * If the data will be read, copy it into temp.
+             * TODO: It might be possible to do a view into
+             *       op[iop]'s mask instead here.
+             */
             if (op_itflags[iop] & NPY_OP_ITFLAG_READ) {
                 if (PyArray_CopyInto(temp, op[iop]) != 0) {
                     Py_DECREF(temp);