ENH: missingdata: Make ndarray.view validate the maskna= parameter be…

…tter
numpy · Aug 27, 2011 · d7076dc · d7076dc
1 parent 0e1a4e9
commit d7076dc
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 4 deletions.
diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c
@@ -171,17 +171,38 @@ array_view(PyArrayObject *self, PyObject *args, PyObject *kwds)
     PyObject *out_type = NULL;
     PyArray_Descr *dtype = NULL;
     PyObject *ret;
-    int maskna = 0, ownmaskna = 0;
+    int maskna = -1, ownmaskna = 0;
+    PyObject *maskna_in = Py_None;
 
     static char *kwlist[] = {"dtype", "type", "maskna", "ownmaskna", NULL};
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OOii", kwlist,
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|OOOi", kwlist,
                                      &out_dtype,
                                      &out_type,
-                                     &maskna,
+                                     &maskna_in,
                                      &ownmaskna)) {
         return NULL;
     }
 
+    /* Treat None the same as not providing the parameter */
+    if (maskna_in != Py_None) {
+        maskna = PyObject_IsTrue(maskna_in);
+        if (maskna == -1) {
+            return NULL;
+        }
+    }
+
+    /* 'ownmaskna' forces 'maskna' to be True */
+    if (ownmaskna) {
+        if (maskna == 0) {
+            PyErr_SetString(PyExc_ValueError,
+                    "cannot specify maskna=False and ownmaskna=True");
+            return NULL;
+        }
+        else {
+            maskna = 1;
+        }
+    }
+
     /* If user specified a positional argument, guess whether it
        represents a type or a dtype for backward compatibility. */
     if (out_dtype) {
@@ -213,8 +234,11 @@ array_view(PyArrayObject *self, PyObject *args, PyObject *kwds)
     }
 
     ret = PyArray_View(self, dtype, (PyTypeObject*)out_type);
+    if (ret == NULL) {
+        return NULL;
+    }
 
-    if (maskna || ownmaskna) {
+    if (maskna == 1) {
         /* Ensure there is an NA mask if requested */
         if (PyArray_AllocateMaskNA((PyArrayObject *)ret,
                                         ownmaskna, 0, 1) < 0) {
@@ -223,6 +247,13 @@ array_view(PyArrayObject *self, PyObject *args, PyObject *kwds)
         }
         return ret;
     }
+    else if (maskna == 0 && PyArray_HASMASKNA(ret)) {
+        PyErr_SetString(PyExc_ValueError,
+                    "Cannot take a view of an NA-masked array "
+                    "with maskna=False");
+        Py_DECREF(ret);
+        return NULL;
+    }
     else {
         return ret;
     }

diff --git a/numpy/core/tests/test_maskna.py b/numpy/core/tests/test_maskna.py
@@ -334,6 +334,9 @@ def test_array_maskna_view_function():
     assert_(c.flags.maskna)
     assert_(not c.flags.ownmaskna)
 
+    # Taking a view of a masked array with maskna=False is invalid
+    assert_raises(ValueError, b.view, maskna=False)
+
     # Taking a view of a masked array, making sure there's a mask
     c = b.view(maskna = True)
     assert_(b.flags.maskna)