ENH: missingdata: In progress making slicing work with NA masks

4 files changed, 88 insertions, 2 deletions

diff --git a/numpy/core/src/multiarray/iterators.c b/numpy/core/src/multiarray/iterators.c
index e1d44740f..d14e1c036 100644
--- a/numpy/core/src/multiarray/iterators.c
+++ b/numpy/core/src/multiarray/iterators.c
@@ -97,7 +97,7 @@ parse_index(PyArrayObject *self, PyObject *op,
     else {
         if (!PySequence_Check(op)) {
             PyErr_SetString(PyExc_IndexError,
-                            "index must be either an int "\
+                            "index must be either an int "
                             "or a sequence");
             return -1;
         }
diff --git a/numpy/core/src/multiarray/mapping.c b/numpy/core/src/multiarray/mapping.c
index 98a8c00f4..0b6ee20d4 100644
--- a/numpy/core/src/multiarray/mapping.c
+++ b/numpy/core/src/multiarray/mapping.c
@@ -595,7 +595,23 @@ array_subscript_simple(PyArrayObject *self, PyObject *op)
     PyArrayObject *other;
     npy_intp value;
 
+    /*
+     * 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(op);
+    if (ind != NULL) {
+        value = PyArray_PyIntAsIntp(ind);
+        Py_DECREF(ind);
+    }
+    else {
+        value = -1;
+    }
+#else
     value = PyArray_PyIntAsIntp(op);
+#endif
     if (value == -1 && PyErr_Occurred()) {
         PyErr_Clear();
     }
diff --git a/numpy/core/src/multiarray/sequence.c b/numpy/core/src/multiarray/sequence.c
index 354dcfa2f..8c286d84d 100644
--- a/numpy/core/src/multiarray/sequence.c
+++ b/numpy/core/src/multiarray/sequence.c
@@ -82,6 +82,28 @@ array_slice(PyArrayObject *self, Py_ssize_t ilow, Py_ssize_t ihigh)
         return NULL;
     }
     PyArray_UpdateFlags(ret, NPY_ARRAY_UPDATE_ALL);
+
+    /* Also take a view of the NA mask if it exists */
+    if (PyArray_HASMASKNA(self)) {
+        PyArrayObject_fieldaccess *fret = (PyArrayObject_fieldaccess *)ret;
+
+        fret->maskna_dtype = PyArray_MASKNA_DTYPE(self);
+        Py_INCREF(fret->maskna_dtype);
+
+        data = PyArray_MASKNA_DATA(self);
+        if (ilow < ihigh) {
+            data += ilow * PyArray_MASKNA_STRIDES(self)[0];
+        }
+        fret->maskna_data = data;
+
+        memcpy(fret->maskna_strides, PyArray_MASKNA_STRIDES(self),
+                        PyArray_NDIM(self) * sizeof(npy_intp));
+
+        /* This view doesn't own the mask */
+        fret->flags |= NPY_ARRAY_MASKNA;
+        fret->flags &= ~NPY_ARRAY_OWNMASKNA;
+    }
+
     return (PyObject *)ret;
 }
 
@@ -102,7 +124,8 @@ array_ass_slice(PyArrayObject *self, Py_ssize_t ilow,
                         "array is not writeable");
         return -1;
     }
-    if ((tmp = (PyArrayObject *)array_slice(self, ilow, ihigh)) == NULL) {
+    tmp = (PyArrayObject *)array_slice(self, ilow, ihigh);
+    if (tmp == NULL) {
         return -1;
     }
     ret = PyArray_CopyObject(tmp, v);
diff --git a/numpy/core/tests/test_na.py b/numpy/core/tests/test_na.py
index d1011ef27..a6e455c4d 100644
--- a/numpy/core/tests/test_na.py
+++ b/numpy/core/tests/test_na.py
@@ -180,5 +180,52 @@ def test_array_maskna_construction():
     assert_equal(a.dtype, np.dtype('O'))
     assert_equal(type(a[2]), np.NAType)
 
+def test_isna():
+    # Objects which are not np.NA or ndarray all return False
+    assert_equal(np.isna(True), False)
+    assert_equal(np.isna("abc"), False)
+    assert_equal(np.isna([1,2,3]), False)
+    assert_equal(np.isna({3:5}), False)
+    # Various NA values return True
+    assert_equal(np.isna(np.NA), True)
+    assert_equal(np.isna(np.NA()), True)
+    assert_equal(np.isna(np.NA(5)), True)
+    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
+    a = np.arange(10)
+
+    # With no mask, it returns all False
+    assert_equal(np.isna(a), False)
+    assert_equal(np.isna(a).shape, (10,))
+
+    # With a mask but no NAs, it still returns all False
+    a.flags.maskna = True
+    assert_equal(np.isna(a), False)
+    assert_equal(np.isna(a).shape, (10,))
+
+    # Checking isna of a single value
+    assert_equal(np.isna(a[4]), False)
+    # Assigning NA to a single value
+    a[3] = np.NA
+    assert_equal(np.isna(a), [0,0,0,1,0,0,0,0,0,0])
+    # Checking isna of a single value
+    assert_equal(np.isna(a[3]), True)
+
+    # Checking isna of a slice
+    assert_equal(np.isna(a[1:6]), [0,0,1,0,0])
+    # Assigning NA to a slice
+    a[5:7] = np.NA
+    assert_equal(np.isna(a), [0,0,0,1,0,1,1,0,0,0])
+
+    # Checking isna of a strided slice
+    assert_equal(np.isna(a[1:8:2]), [0,1,1,0])
+    # Assigning NA to a strided slice
+    a[2:10:3] = np.NA
+    assert_equal(np.isna(a), [0,0,1,1,0,1,1,0,1,0])
+
+
 if __name__ == "__main__":
     run_module_suite()