ENH: missingdata: Future-proof AssignNA and AssignMaskNA for later multi-NA support

7 files changed, 39 insertions, 30 deletions

diff --git a/numpy/core/src/multiarray/array_assign_array.c b/numpy/core/src/multiarray/array_assign_array.c
index eae1d10d1..ddad6a191 100644
--- a/numpy/core/src/multiarray/array_assign_array.c
+++ b/numpy/core/src/multiarray/array_assign_array.c
@@ -404,7 +404,8 @@ PyArray_AssignArray(PyArrayObject *dst, PyArrayObject *src,
 
             if (na != NULL) {
                 /* TODO: With multi-NA, preservena must also be followed */
-                int retcode = PyArray_AssignNA(dst, wheremask, na);
+                int retcode = PyArray_AssignNA(dst, na, wheremask,
+                                            preservena, preservewhichna);
                 Py_DECREF(na);
                 return retcode;
             }
@@ -608,7 +609,8 @@ PyArray_AssignArray(PyArrayObject *dst, PyArrayObject *src,
                     goto finish;
                 }
                 else {
-                    if (PyArray_AssignMaskNA(dst, NULL, 1) < 0) {
+                    if (PyArray_AssignMaskNA(dst, 1, NULL,
+                                        preservena, preservewhichna) < 0) {
                         goto fail;
                     }
                 }
@@ -733,7 +735,8 @@ PyArray_AssignArray(PyArrayObject *dst, PyArrayObject *src,
                     goto finish;
                 }
                 else {
-                    if (PyArray_AssignMaskNA(dst, wheremask, 1) < 0) {
+                    if (PyArray_AssignMaskNA(dst, 1, wheremask,
+                                        preservena, preservewhichna) < 0) {
                         goto fail;
                     }
                 }
diff --git a/numpy/core/src/multiarray/array_assign_scalar.c b/numpy/core/src/multiarray/array_assign_scalar.c
index a1e2192c1..2dbad89e4 100644
--- a/numpy/core/src/multiarray/array_assign_scalar.c
+++ b/numpy/core/src/multiarray/array_assign_scalar.c
@@ -405,7 +405,8 @@ PyArray_AssignRawScalar(PyArrayObject *dst,
         if (!preservena || !dst_has_maskna) {
             /* If assigning to an array with an NA mask, set to all exposed */
             if (dst_has_maskna) {
-                if (PyArray_AssignMaskNA(dst, NULL, 1) < 0) {
+                if (PyArray_AssignMaskNA(dst, 1, NULL,
+                                    preservena, preservewhichna) < 0) {
                     goto fail;
                 }
             }
@@ -468,7 +469,8 @@ PyArray_AssignRawScalar(PyArrayObject *dst,
                  * TODO: If the where mask has NA values, this part
                  *       changes too.
                  */
-                if (PyArray_AssignMaskNA(dst, wheremask, 1) < 0) {
+                if (PyArray_AssignMaskNA(dst, 1, wheremask,
+                                    preservena, preservewhichna) < 0) {
                     goto fail;
                 }
             }
diff --git a/numpy/core/src/multiarray/arrayobject.c b/numpy/core/src/multiarray/arrayobject.c
index d21a78d74..f24ea376e 100644
--- a/numpy/core/src/multiarray/arrayobject.c
+++ b/numpy/core/src/multiarray/arrayobject.c
@@ -210,7 +210,7 @@ PyArray_CopyObject(PyArrayObject *dest, PyObject *src_object)
             }
             /* Assigning NA affects the mask if it exists */
             else if (na != NULL) {
-                if (PyArray_AssignNA(dest, NULL, na) < 0) {
+                if (PyArray_AssignNA(dest, na, NULL, 0, NULL) < 0) {
                     Py_DECREF(na);
                     Py_DECREF(src_object);
                     return -1;
diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
index 183f90e9b..ea1c94008 100644
--- a/numpy/core/src/multiarray/ctors.c
+++ b/numpy/core/src/multiarray/ctors.c
@@ -2638,7 +2638,7 @@ PyArray_CopyAsFlat(PyArrayObject *dst, PyArrayObject *src, NPY_ORDER order)
      * to be exposed, then proceed without worrying about the mask.
      */
     else if (PyArray_HASMASKNA(dst)) {
-        if (PyArray_AssignMaskNA(dst, NULL, 1) < 0) {
+        if (PyArray_AssignMaskNA(dst, 1, NULL, 0, NULL) < 0) {
             return -1;
         }
         baseflags |= NPY_ITER_IGNORE_MASKNA;
diff --git a/numpy/core/src/multiarray/na_mask.c b/numpy/core/src/multiarray/na_mask.c
index 2684c472c..03a9dfce0 100644
--- a/numpy/core/src/multiarray/na_mask.c
+++ b/numpy/core/src/multiarray/na_mask.c
@@ -228,11 +228,18 @@ fill_raw_byte_array(int ndim, npy_intp *shape,
  * If 'wheremask' isn't NULL, it should be a boolean mask which
  * specifies where to do the assignment.
  *
+ * The parameters 'preservena' and 'preservewhichna' are NOT YET
+ * SUPPORTED, but are in place to allow for future expansion to
+ * multi-NA. 'preservewhichna' should be set to NULL, while
+ * preservena has no effect for straight NPY_BOOL NA masks, because
+ * different NAs are indistinguishable.
+ *
  * Returns 0 on success, -1 on failure.
  */
 NPY_NO_EXPORT int
-PyArray_AssignMaskNA(PyArrayObject *arr, PyArrayObject *wheremask,
-                        npy_mask maskvalue)
+PyArray_AssignMaskNA(PyArrayObject *arr, npy_mask maskvalue,
+                PyArrayObject *wheremask,
+                npy_bool preservena, npy_bool *preservewhichna)
 {
     PyArray_Descr *maskvalue_dtype;
     int retcode = 0;
@@ -245,6 +252,12 @@ PyArray_AssignMaskNA(PyArrayObject *arr, PyArrayObject *wheremask,
         return -1;
     }
 
+    if (preservewhichna != NULL) {
+        PyErr_SetString(PyExc_RuntimeError,
+                "multi-NA support is not yet implemented");
+        return -1;
+    }
+
     /*
      * If the mask given has no payload, assign from boolean type, otherwise
      * assign from the mask type.
@@ -440,21 +453,22 @@ PyArray_AllocateMaskNA(PyArrayObject *arr,
  * In the future, when 'arr' has an NA dtype, will assign the
  * appropriate NA bitpatterns to the elements.
  *
+ * The parameters 'preservena' and 'preservewhichna' are NOT YET
+ * SUPPORTED, but are in place to allow for future expansion to
+ * multi-NA. 'preservewhichna' should be set to NULL, while
+ * preservena has no effect for straight NPY_BOOL NA masks, because
+ * different NAs are indistinguishable.
+ *
  * Returns -1 on failure, 0 on success.
  */
 NPY_NO_EXPORT int
-PyArray_AssignNA(PyArrayObject *arr, PyArrayObject *wheremask, NpyNA *na)
+PyArray_AssignNA(PyArrayObject *arr, NpyNA *na,
+                PyArrayObject *wheremask,
+                npy_bool preservena, npy_bool *preservewhichna)
 {
     NpyNA_fields *fna = (NpyNA_fields *)na;
     char maskvalue;
 
-    if (!PyArray_HASMASKNA(arr)) {
-        PyErr_SetString(PyExc_ValueError,
-                "Cannot assign an NA to an "
-                "array with no NA support");
-        return -1;
-    }
-
     /* Turn the payload into a mask value */
     if (fna->payload == NPY_NA_NOPAYLOAD) {
         maskvalue = 0;
@@ -471,7 +485,8 @@ PyArray_AssignNA(PyArrayObject *arr, PyArrayObject *wheremask, NpyNA *na)
         maskvalue = (char)NpyMaskValue_Create(0, fna->payload);
     }
 
-    return PyArray_AssignMaskNA(arr, wheremask, maskvalue);
+    return PyArray_AssignMaskNA(arr, maskvalue,
+                        wheremask, preservena, preservewhichna);
 }
 
 /*
diff --git a/numpy/core/src/multiarray/na_mask.h b/numpy/core/src/multiarray/na_mask.h
index 547389e77..160327de2 100644
--- a/numpy/core/src/multiarray/na_mask.h
+++ b/numpy/core/src/multiarray/na_mask.h
@@ -4,17 +4,6 @@
 #include "lowlevel_strided_loops.h"
 
 /*
- * Assigns the given NA value to all the elements in the array.
- *
- * If 'wheremask' isn't NULL, it specifies which elements to assign
- * NA to.
- *
- * Returns -1 on failure, 0 on success.
- */
-NPY_NO_EXPORT int
-PyArray_AssignNA(PyArrayObject *arr, PyArrayObject *wheremask, NpyNA *na);
-
-/*
  * A ufunc-like function, which returns a boolean or an array
  * of booleans indicating which values are NA.
  */
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index c1ab4055d..9de0e7ef5 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -1010,7 +1010,7 @@ static int get_ufunc_arguments(PyUFuncObject *ufunc,
     else if (!(*out_use_maskna) && any_maskna_out) {
         for (i = nin; i < nin+nout; ++i) {
             if (PyArray_HASMASKNA(out_op[i])) {
-                if (PyArray_AssignMaskNA(out_op[i], NULL, 1) < 0) {
+                if (PyArray_AssignMaskNA(out_op[i], 1, NULL, 0, NULL) < 0) {
                     return -1;
                 }
             }