ENH: missingdata: Get the NA mask working with slice indexing

4 files changed, 182 insertions, 94 deletions

diff --git a/numpy/core/src/multiarray/iterators.c b/numpy/core/src/multiarray/iterators.c
index d14e1c036..8371702f9 100644
--- a/numpy/core/src/multiarray/iterators.c
+++ b/numpy/core/src/multiarray/iterators.c
@@ -17,22 +17,31 @@
 #include "ctors.h"
 #include "common.h"
 
-#define PseudoIndex -1
-#define RubberIndex -2
-#define SingleIndex -3
+#define NEWAXIS_INDEX -1
+#define ELLIPSIS_INDEX -2
+#define SINGLE_INDEX -3
 
+static int
+slice_coerce_index(PyObject *o, npy_intp *v);
+
+/*
+ * This function converts one element of the indexing tuple
+ * into a step size and a number of steps, returning the
+ * starting index. Non-slices are signalled in 'n_steps',
+ * as NEWAXIS_INDEX, ELLIPSIS_INDEX, or SINGLE_INDEX.
+ */
 NPY_NO_EXPORT npy_intp
-parse_subindex(PyObject *op, npy_intp *step_size,
+parse_index_entry(PyObject *op, npy_intp *step_size,
                     npy_intp *n_steps, npy_intp max)
 {
     npy_intp index;
 
     if (op == Py_None) {
-        *n_steps = PseudoIndex;
+        *n_steps = NEWAXIS_INDEX;
         index = 0;
     }
     else if (op == Py_Ellipsis) {
-        *n_steps = RubberIndex;
+        *n_steps = ELLIPSIS_INDEX;
         index = 0;
     }
     else if (PySlice_Check(op)) {
@@ -52,15 +61,14 @@ parse_subindex(PyObject *op, npy_intp *step_size,
         }
     }
     else {
-        index = PyArray_PyIntAsIntp(op);
-        if (error_converting(index)) {
+        if (!slice_coerce_index(op, &index)) {
             PyErr_SetString(PyExc_IndexError,
-                            "each subindex must be either a "\
-                            "slice, an integer, Ellipsis, or "\
+                            "each index entry must be either a "
+                            "slice, an integer, Ellipsis, or "
                             "newaxis");
             goto fail;
         }
-        *n_steps = SingleIndex;
+        *n_steps = SINGLE_INDEX;
         *step_size = 0;
         if (index < 0) {
             index += max;
@@ -77,13 +85,23 @@ parse_subindex(PyObject *op, npy_intp *step_size,
 }
 
 
+/*
+ * Parses an index that has no fancy indexing. Populates
+ * out_dimensions, out_strides, and out_offset. If out_maskna_strides
+ * and out_maskoffset aren't NULL, then 'self' must have an NA mask
+ * which is used to populate those variables as well.
+ */
 NPY_NO_EXPORT int
 parse_index(PyArrayObject *self, PyObject *op,
-            npy_intp *dimensions, npy_intp *strides, npy_intp *offset_ptr)
+            npy_intp *out_dimensions,
+            npy_intp *out_strides,
+            npy_intp *out_offset,
+            npy_intp *out_maskna_strides,
+            npy_intp *out_maskna_offset)
 {
     int i, j, n;
-    int nd_old, nd_new, n_add, n_pseudo;
-    npy_intp n_steps, start, offset, step_size;
+    int nd_old, nd_new, n_add, n_ellipsis;
+    npy_intp n_steps, start, offset, maskna_offset, step_size;
     PyObject *op1 = NULL;
     int is_slice;
 
@@ -108,62 +126,71 @@ parse_index(PyArrayObject *self, PyObject *op,
     nd_old = nd_new = 0;
 
     offset = 0;
+    maskna_offset = 0;
     for (i = 0; i < n; i++) {
         if (!is_slice) {
-            if (!(op1=PySequence_GetItem(op, i))) {
-                PyErr_SetString(PyExc_IndexError,
-                                "invalid index");
+            op1 = PySequence_GetItem(op, i);
+            if (op1 == NULL) {
                 return -1;
             }
         }
-        start = parse_subindex(op1, &step_size, &n_steps,
+        start = parse_index_entry(op1, &step_size, &n_steps,
                                nd_old < PyArray_NDIM(self) ?
                                PyArray_DIMS(self)[nd_old] : 0);
         Py_DECREF(op1);
         if (start == -1) {
             break;
         }
-        if (n_steps == PseudoIndex) {
-            dimensions[nd_new] = 1; strides[nd_new] = 0;
+        if (n_steps == NEWAXIS_INDEX) {
+            out_dimensions[nd_new] = 1;
+            out_strides[nd_new] = 0;
+            if (out_maskna_strides != NULL) {
+                out_maskna_strides[nd_new] = 0;
+            }
             nd_new++;
         }
-        else {
-            if (n_steps == RubberIndex) {
-                for (j = i + 1, n_pseudo = 0; j < n; j++) {
-                    op1 = PySequence_GetItem(op, j);
-                    if (op1 == Py_None) {
-                        n_pseudo++;
-                    }
-                    Py_DECREF(op1);
-                }
-                n_add = PyArray_NDIM(self)-(n-i-n_pseudo-1+nd_old);
-                if (n_add < 0) {
-                    PyErr_SetString(PyExc_IndexError,
-                                    "too many indices");
-                    return -1;
-                }
-                for (j = 0; j < n_add; j++) {
-                    dimensions[nd_new] = \
-                        PyArray_DIMS(self)[nd_old];
-                    strides[nd_new] = \
-                        PyArray_STRIDES(self)[nd_old];
-                    nd_new++; nd_old++;
+        else if (n_steps == ELLIPSIS_INDEX) {
+            for (j = i + 1, n_ellipsis = 0; j < n; j++) {
+                op1 = PySequence_GetItem(op, j);
+                if (op1 == Py_None) {
+                    n_ellipsis++;
                 }
+                Py_DECREF(op1);
             }
-            else {
-                if (nd_old >= PyArray_NDIM(self)) {
-                    PyErr_SetString(PyExc_IndexError,
-                                    "too many indices");
-                    return -1;
+            n_add = PyArray_NDIM(self)-(n-i-n_ellipsis-1+nd_old);
+            if (n_add < 0) {
+                PyErr_SetString(PyExc_IndexError, "too many indices");
+                return -1;
+            }
+            for (j = 0; j < n_add; j++) {
+                out_dimensions[nd_new] = PyArray_DIMS(self)[nd_old];
+                out_strides[nd_new] = PyArray_STRIDES(self)[nd_old];
+                if (out_maskna_strides != NULL) {
+                    out_maskna_strides[nd_new] =
+                                    PyArray_MASKNA_STRIDES(self)[nd_old];
                 }
-                offset += PyArray_STRIDES(self)[nd_old]*start;
-                nd_old++;
-                if (n_steps != SingleIndex) {
-                    dimensions[nd_new] = n_steps;
-                    strides[nd_new] = step_size * \
-                        PyArray_STRIDES(self)[nd_old-1];
-                    nd_new++;
+                nd_new++; nd_old++;
+            }
+        }
+        else {
+            if (nd_old >= PyArray_NDIM(self)) {
+                PyErr_SetString(PyExc_IndexError, "too many indices");
+                return -1;
+            }
+            offset += PyArray_STRIDES(self)[nd_old]*start;
+            if (out_maskna_offset != NULL) {
+                maskna_offset += PyArray_MASKNA_STRIDES(self)[nd_old]*start;
+            }
+            nd_old++;
+            if (n_steps != SINGLE_INDEX) {
+                out_dimensions[nd_new] = n_steps;
+                out_strides[nd_new] = step_size *
+                                            PyArray_STRIDES(self)[nd_old-1];
+                if (out_maskna_strides != NULL) {
+                    out_maskna_strides[nd_new] = step_size *
+                                        PyArray_MASKNA_STRIDES(self)[nd_old-1];
                 }
+                nd_new++;
             }
         }
     }
@@ -172,20 +199,47 @@ parse_index(PyArrayObject *self, PyObject *op,
     }
     n_add = PyArray_NDIM(self)-nd_old;
     for (j = 0; j < n_add; j++) {
-        dimensions[nd_new] = PyArray_DIMS(self)[nd_old];
-        strides[nd_new] = PyArray_STRIDES(self)[nd_old];
+        out_dimensions[nd_new] = PyArray_DIMS(self)[nd_old];
+        out_strides[nd_new] = PyArray_STRIDES(self)[nd_old];
+        if (out_maskna_strides != NULL) {
+            out_maskna_strides[nd_new] = PyArray_MASKNA_STRIDES(self)[nd_old];
+        }
         nd_new++;
         nd_old++;
     }
-    *offset_ptr = offset;
+    *out_offset = offset;
+    if (out_maskna_offset != NULL) {
+        *out_maskna_offset = maskna_offset;
+    }
     return nd_new;
 }
 
+/*
+ * Tries to convert 'o' into an npy_intp interpreted as an
+ * index. Returns 1 if it was successful, 0 otherwise. Does
+ * not set an exception.
+ */
 static int
 slice_coerce_index(PyObject *o, npy_intp *v)
 {
+    /*
+     * PyNumber_Index was introduced in Python 2.5 because of NumPy.
+     * http://www.python.org/dev/peps/pep-0357/
+     * Let's use it for indexing!
+     */
+#if PY_VERSION_HEX >= 0x02050000
+    PyObject *ind = PyNumber_Index(o);
+    if (ind != NULL) {
+        *v = PyArray_PyIntAsIntp(ind);
+        Py_DECREF(ind);
+    }
+    else {
+        *v = -1;
+    }
+#else
     *v = PyArray_PyIntAsIntp(o);
-    if (error_converting(*v)) {
+#endif
+    if ((*v) == -1 && PyErr_Occurred()) {
         PyErr_Clear();
         return 0;
     }
@@ -769,18 +823,18 @@ iter_subscript(PyArrayIterObject *self, PyObject *ind)
 
     /* Check for Integer or Slice */
     if (PyLong_Check(ind) || PyInt_Check(ind) || PySlice_Check(ind)) {
-        start = parse_subindex(ind, &step_size, &n_steps,
+        start = parse_index_entry(ind, &step_size, &n_steps,
                                self->size);
         if (start == -1) {
             goto fail;
         }
-        if (n_steps == RubberIndex || n_steps == PseudoIndex) {
+        if (n_steps == ELLIPSIS_INDEX || n_steps == NEWAXIS_INDEX) {
             PyErr_SetString(PyExc_IndexError,
                             "cannot use Ellipsis or newaxes here");
             goto fail;
         }
         PyArray_ITER_GOTO1D(self, start)
-            if (n_steps == SingleIndex) { /* Integer */
+            if (n_steps == SINGLE_INDEX) { /* Integer */
                 PyObject *tmp;
                 tmp = PyArray_ToScalar(self->dataptr, self->ao);
                 PyArray_ITER_RESET(self);
@@ -1041,17 +1095,17 @@ iter_ass_subscript(PyArrayIterObject *self, PyObject *ind, PyObject *val)
 
     /* Check Slice */
     if (PySlice_Check(ind)) {
-        start = parse_subindex(ind, &step_size, &n_steps, self->size);
+        start = parse_index_entry(ind, &step_size, &n_steps, self->size);
         if (start == -1) {
             goto finish;
         }
-        if (n_steps == RubberIndex || n_steps == PseudoIndex) {
+        if (n_steps == ELLIPSIS_INDEX || n_steps == NEWAXIS_INDEX) {
             PyErr_SetString(PyExc_IndexError,
                             "cannot use Ellipsis or newaxes here");
             goto finish;
         }
         PyArray_ITER_GOTO1D(self, start);
-        if (n_steps == SingleIndex) {
+        if (n_steps == SINGLE_INDEX) {
             /* Integer */
             copyswap(self->dataptr, PyArray_DATA(arrval), swap, arrval);
             PyArray_ITER_RESET(self);
diff --git a/numpy/core/src/multiarray/iterators.h b/numpy/core/src/multiarray/iterators.h
index 3099425c5..e877f8520 100644
--- a/numpy/core/src/multiarray/iterators.h
+++ b/numpy/core/src/multiarray/iterators.h
@@ -1,12 +1,19 @@
 #ifndef _NPY_ARRAYITERATORS_H_
 #define _NPY_ARRAYITERATORS_H_
 
-NPY_NO_EXPORT intp
-parse_subindex(PyObject *op, intp *step_size, intp *n_steps, intp max);
-
+/*
+ * Parses an index that has no fancy indexing. Populates
+ * out_dimensions, out_strides, and out_offset. If out_maskstrides
+ * and out_maskoffset aren't NULL, then 'self' must have an NA mask
+ * which is used to populate those variables as well.
+ */
 NPY_NO_EXPORT int
 parse_index(PyArrayObject *self, PyObject *op,
-            intp *dimensions, intp *strides, intp *offset_ptr);
+            npy_intp *out_dimensions,
+            npy_intp *out_strides,
+            npy_intp *out_offset,
+            npy_intp *out_maskna_strides,
+            npy_intp *out_maskna_offset);
 
 NPY_NO_EXPORT PyObject
 *iter_subscript(PyArrayIterObject *, PyObject *);
diff --git a/numpy/core/src/multiarray/mapping.c b/numpy/core/src/multiarray/mapping.c
index 0b6ee20d4..396f0a9e4 100644
--- a/numpy/core/src/multiarray/mapping.c
+++ b/numpy/core/src/multiarray/mapping.c
@@ -222,7 +222,7 @@ _swap_axes(PyArrayMapIterObject *mit, PyArrayObject **ret, int getmap)
     int n1, n2, n3, val, bnd;
     int i;
     PyArray_Dims permute;
-    npy_intp d[MAX_DIMS];
+    npy_intp d[NPY_MAXDIMS];
     PyArrayObject *arr;
 
     permute.ptr = d;
@@ -454,7 +454,7 @@ count_new_axes_0d(PyObject *tuple)
                         " as an index");
         return -1;
     }
-    if (newaxis_count > MAX_DIMS) {
+    if (newaxis_count > NPY_MAXDIMS) {
         PyErr_SetString(PyExc_IndexError, "too many dimensions");
         return -1;
     }
@@ -465,7 +465,7 @@ NPY_NO_EXPORT PyObject *
 add_new_axes_0d(PyArrayObject *arr,  int newaxis_count)
 {
     PyArrayObject *other;
-    npy_intp dimensions[MAX_DIMS];
+    npy_intp dimensions[NPY_MAXDIMS];
     int i;
 
     for (i = 0; i < newaxis_count; ++i) {
@@ -501,7 +501,7 @@ fancy_indexing_check(PyObject *args)
 
     if (PyTuple_Check(args)) {
         n = PyTuple_GET_SIZE(args);
-        if (n >= MAX_DIMS) {
+        if (n >= NPY_MAXDIMS) {
             return SOBJ_TOOMANY;
         }
         for (i = 0; i < n; i++) {
@@ -532,14 +532,14 @@ fancy_indexing_check(PyObject *args)
     }
     else if (PySequence_Check(args)) {
         /*
-         * Sequences < MAX_DIMS with any slice objects
+         * Sequences < NPY_MAXDIMS with any slice objects
          * or newaxis, or Ellipsis is considered standard
          * as long as there are also no Arrays and or additional
          * sequences embedded.
          */
         retval = SOBJ_ISFANCY;
         n = PySequence_Size(args);
-        if (n < 0 || n >= MAX_DIMS) {
+        if (n < 0 || n >= NPY_MAXDIMS) {
             return SOBJ_ISFANCY;
         }
         for (i = 0; i < n; i++) {
@@ -589,10 +589,11 @@ fancy_indexing_check(PyObject *args)
 NPY_NO_EXPORT PyObject *
 array_subscript_simple(PyArrayObject *self, PyObject *op)
 {
-    npy_intp dimensions[MAX_DIMS], strides[MAX_DIMS];
-    npy_intp offset;
+    npy_intp dimensions[NPY_MAXDIMS], strides[NPY_MAXDIMS];
+    npy_intp maskna_strides[NPY_MAXDIMS];
+    npy_intp offset, maskna_offset;
     int nd;
-    PyArrayObject *other;
+    PyArrayObject *ret;
     npy_intp value;
 
     /*
@@ -620,30 +621,55 @@ array_subscript_simple(PyArrayObject *self, PyObject *op)
     }
 
     /* Standard (view-based) Indexing */
-    nd = parse_index(self, op, dimensions, strides, &offset);
+    if (PyArray_HASMASKNA(self)) {
+        nd = parse_index(self, op, dimensions,
+                        strides, &offset, maskna_strides, &maskna_offset);
+    }
+    else {
+        nd = parse_index(self, op, dimensions,
+                        strides, &offset, NULL, NULL);
+    }
     if (nd == -1) {
         return NULL;
     }
 
-    /* This will only work if new array will be a view */
+    /* Create a view using the indexing result */
     Py_INCREF(PyArray_DESCR(self));
-    other = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(self),
+    ret = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(self),
                                 PyArray_DESCR(self),
                                 nd, dimensions,
-                                strides, PyArray_DATA(self)+offset,
+                                strides, PyArray_DATA(self) + offset,
                                 PyArray_FLAGS(self),
                                 (PyObject *)self);
-    if (other == NULL) {
+    if (ret == NULL) {
         return NULL;
     }
     Py_INCREF(self);
-    if (PyArray_SetBaseObject(other, (PyObject *)self) < 0) {
-        Py_DECREF(other);
+    if (PyArray_SetBaseObject(ret, (PyObject *)self) < 0) {
+        Py_DECREF(ret);
         return NULL;
     }
-    PyArray_UpdateFlags(other, NPY_ARRAY_UPDATE_ALL);
+    PyArray_UpdateFlags(ret, NPY_ARRAY_UPDATE_ALL);
 
-    return (PyObject *)other;
+    if (PyArray_HASMASKNA(self)) {
+        PyArrayObject_fieldaccess *fret = (PyArrayObject_fieldaccess *)ret;
+
+        fret->maskna_dtype = PyArray_MASKNA_DTYPE(self);
+        Py_INCREF(fret->maskna_dtype);
+
+        fret->maskna_data = PyArray_MASKNA_DATA(self) + maskna_offset;
+
+        if (nd > 0) {
+            memcpy(fret->maskna_strides, maskna_strides,
+                                            nd * sizeof(npy_intp));
+        }
+
+        /* This view doesn't own the mask */
+        fret->flags |= NPY_ARRAY_MASKNA;
+        fret->flags &= ~NPY_ARRAY_OWNMASKNA;
+    }
+
+    return (PyObject *)ret;
 }
 
 NPY_NO_EXPORT PyObject *
@@ -884,7 +910,7 @@ array_ass_sub(PyArrayObject *self, PyObject *index, PyObject *op)
 {
     int ret, oned, fancy;
     PyArrayMapIterObject *mit;
-    npy_intp vals[MAX_DIMS];
+    npy_intp vals[NPY_MAXDIMS];
 
     if (op == NULL) {
         PyErr_SetString(PyExc_ValueError,
@@ -1048,7 +1074,7 @@ array_subscript_nice(PyArrayObject *self, PyObject *op)
 {
 
     PyArrayObject *mp;
-    npy_intp vals[MAX_DIMS];
+    npy_intp vals[NPY_MAXDIMS];
 
     if (PyInt_Check(op) || PyArray_IsScalar(op, Integer) ||
         PyLong_Check(op) || (PyIndex_Check(op) &&
@@ -1175,8 +1201,8 @@ _nonzero_indices(PyObject *myBool, PyArrayIterObject **iters)
     int nd, j;
     npy_intp size, i, count;
     Bool *ptr;
-    npy_intp coords[MAX_DIMS], dims_m1[MAX_DIMS];
-    npy_intp *dptr[MAX_DIMS];
+    npy_intp coords[NPY_MAXDIMS], dims_m1[NPY_MAXDIMS];
+    npy_intp *dptr[NPY_MAXDIMS];
 
     typecode=PyArray_DescrFromType(NPY_BOOL);
     ba = (PyArrayObject *)PyArray_FromAny(myBool, typecode, 0, 0,
@@ -1292,7 +1318,7 @@ _convert_obj(PyObject *obj, PyArrayIterObject **iter)
 NPY_NO_EXPORT void
 PyArray_MapIterReset(PyArrayMapIterObject *mit)
 {
-    int i,j; npy_intp coord[MAX_DIMS];
+    int i,j; npy_intp coord[NPY_MAXDIMS];
     PyArrayIterObject *it;
     PyArray_CopySwapFunc *copyswap;
 
@@ -1340,7 +1366,7 @@ NPY_NO_EXPORT void
 PyArray_MapIterNext(PyArrayMapIterObject *mit)
 {
     int i, j;
-    npy_intp coord[MAX_DIMS];
+    npy_intp coord[NPY_MAXDIMS];
     PyArrayIterObject *it;
     PyArray_CopySwapFunc *copyswap;
 
@@ -1588,7 +1614,7 @@ PyArray_MapIterNew(PyObject *indexobj, int oned, int fancy)
     if (mit == NULL) {
         return NULL;
     }
-    for (i = 0; i < MAX_DIMS; i++) {
+    for (i = 0; i < NPY_MAXDIMS; i++) {
         mit->iters[i] = NULL;
     }
     mit->index = 0;
diff --git a/numpy/core/tests/test_na.py b/numpy/core/tests/test_na.py
index a6e455c4d..3b179b126 100644
--- a/numpy/core/tests/test_na.py
+++ b/numpy/core/tests/test_na.py
@@ -193,8 +193,7 @@ def test_isna():
     assert_equal(np.isna(np.NA(dtype='f4')), True)
     assert_equal(np.isna(np.NA(12,dtype='f4')), True)
 
-def test_array_maskna_isna():
-    # Simple 1D example
+def test_array_maskna_isna_1D():
     a = np.arange(10)
 
     # With no mask, it returns all False
@@ -226,6 +225,8 @@ def test_array_maskna_isna():
     a[2:10:3] = np.NA
     assert_equal(np.isna(a), [0,0,1,1,0,1,1,0,1,0])
 
+    # TODO: fancy indexing is next...
+
 
 if __name__ == "__main__":
     run_module_suite()