TST: missingdata: Write some tests for the np.any and np.all NA behavior

2 files changed, 87 insertions, 7 deletions

diff --git a/numpy/core/src/multiarray/boolean_ops.c.src b/numpy/core/src/multiarray/boolean_ops.c.src
index da39824f0..f3d87cc34 100644
--- a/numpy/core/src/multiarray/boolean_ops.c.src
+++ b/numpy/core/src/multiarray/boolean_ops.c.src
@@ -29,11 +29,15 @@
 /* Typedef for the reduction inner loop */
 typedef void (reduce_inner_loop)(char **, npy_intp *, npy_intp);
 
+#define ANY_COMBINE_MASK(val, mask) ((val) & (mask))
+#define ALL_COMBINE_MASK(val, mask) ((val) | ((mask) ^ 0x01))
+
 /**begin repeat
  * #oper = any, all#
  * #Oper = Any, All#
  * #Identity = Zero, One#
  * #combineop = |=, &=#
+ * #combinemask = ANY_COMBINE_MASK, ALL_COMBINE_MASK#
  * #shortcircuit = !value, value#
  * #idval = 0, 1#
  */
@@ -164,7 +168,7 @@ static void
     npy_intp i;
 
     for (i = 0; i < count; ++i) {
-        *data0 @combineop@ *data1 & *data2;
+        *data0 @combineop@ @combinemask@(*data1, *data2);
 
         data0 += stride0;
         data1 += stride1;
@@ -182,7 +186,7 @@ static void
     char value = *data0;
 
     for (i = 0; i < count && @shortcircuit@; ++i) {
-        value @combineop@ *data1 & *data2;
+        value @combineop@ @combinemask@(*data1, *data2);
 
         data1 += stride1;
         data2 += stride2;
@@ -281,9 +285,9 @@ static void
          * If the value is an exposed True (for any) or False (for all),
          * expose the result as well
          */
-        else if (*data1 & *data3) {
+        else if (@combinemask@(*data1, *data3) == 1-@idval@) {
             *data0 = 1-@idval@;
-            *data2 = 1-@idval@;
+            *data2 = 1;
         }
 
         data0 += stride0;
@@ -302,7 +306,7 @@ static void
     npy_intp stride1 = strides[1], stride3 = strides[3];
     npy_intp i;
     char maskvalue = *data2;
-    char value = maskvalue ? *data0 : 0;
+    char value = maskvalue ? *data0 : @idval@;
 
     for (i = 0; i < count && @shortcircuit@; ++i) {
         /* Normal case */
@@ -313,9 +317,9 @@ static void
          * If the value is an exposed True (for any) or False (for all),
          * expose the result as well
          */
-        else if (*data1 & *data3) {
+        else if (@combinemask@(*data1, *data3) == 1-@idval@) {
             value = 1-@idval@;
-            maskvalue = 1-@idval@;
+            maskvalue = 1;
             break;
         }
 
diff --git a/numpy/core/tests/test_maskna.py b/numpy/core/tests/test_maskna.py
index f83dab264..bbda9f5fa 100644
--- a/numpy/core/tests/test_maskna.py
+++ b/numpy/core/tests/test_maskna.py
@@ -1023,6 +1023,82 @@ def check_ufunc_skipna_max_3D(max_func):
     assert_raises(ValueError, max_func, a, axis=1, skipna=True)
     assert_raises(ValueError, max_func, a, axis=2, skipna=True)
 
+def test_ufunc_ndarray_any():
+    a = np.array([0,0,0,0,0], dtype='?', maskna=True)
+    assert_array_equal(a.any(), False)
+    a[0] = np.NA
+    assert_array_equal(a.any(), np.NA)
+    assert_array_equal(a.any(skipna=True), False)
+    a[0] = 0
+    a[-1] = np.NA
+    assert_array_equal(a.any(), np.NA)
+    assert_array_equal(a.any(skipna=True), False)
+    a[0] = 1
+    assert_array_equal(a.any(), True)
+    assert_array_equal(a.any(skipna=True), True)
+    a[-1] = 1
+    a[-2] = np.NA
+    assert_array_equal(a.any(), True)
+    assert_array_equal(a.any(skipna=True), True)
+
+    a = np.array([[0,0,0],[0,np.NA,0]], dtype='?')
+    assert_array_equal(a.any(axis=0), [False, np.NA, False])
+    assert_array_equal(a.any(axis=1), [False, np.NA])
+    assert_array_equal(a.any(axis=0, skipna=True), [False, False, False])
+    assert_array_equal(a.any(axis=1, skipna=True), [False, False])
+
+    a[0,1] = 1
+    assert_array_equal(a.any(axis=0), [False, True, False])
+    assert_array_equal(a.any(axis=1), [True, np.NA])
+    assert_array_equal(a.any(axis=0, skipna=True), [False, True, False])
+    assert_array_equal(a.any(axis=1, skipna=True), [True, False])
+
+    a[0,1] = np.NA
+    a[1,1] = 0
+    a[0,2] = 1
+    assert_array_equal(a.any(axis=0), [False, np.NA, True])
+    assert_array_equal(a.any(axis=1), [True, False])
+    assert_array_equal(a.any(axis=0, skipna=True), [False, False, True])
+    assert_array_equal(a.any(axis=1, skipna=True), [True, False])
+
+def test_ufunc_ndarray_all():
+    a = np.array([1,1,1,1,1], dtype='?', maskna=True)
+    assert_array_equal(a.all(), True)
+    a[0] = np.NA
+    assert_array_equal(a.all(), np.NA)
+    assert_array_equal(a.all(skipna=True), True)
+    a[0] = 1
+    a[-1] = np.NA
+    assert_array_equal(a.all(), np.NA)
+    assert_array_equal(a.all(skipna=True), True)
+    a[0] = 0
+    assert_array_equal(a.all(), False)
+    assert_array_equal(a.all(skipna=True), False)
+    a[-1] = 0
+    a[-2] = np.NA
+    assert_array_equal(a.all(), False)
+    assert_array_equal(a.all(skipna=True), False)
+
+    a = np.array([[1,1,1],[1,np.NA,1]], dtype='?')
+    assert_array_equal(a.all(axis=0), [True, np.NA, True])
+    assert_array_equal(a.all(axis=1), [True, np.NA])
+    assert_array_equal(a.all(axis=0, skipna=True), [True, True, True])
+    assert_array_equal(a.all(axis=1, skipna=True), [True, True])
+
+    a[0,1] = 0
+    assert_array_equal(a.all(axis=0), [True, False, True])
+    assert_array_equal(a.all(axis=1), [False, np.NA])
+    assert_array_equal(a.all(axis=0, skipna=True), [True, False, True])
+    assert_array_equal(a.all(axis=1, skipna=True), [False, True])
+
+    a[0,1] = np.NA
+    a[1,1] = 1
+    a[0,2] = 0
+    assert_array_equal(a.all(axis=0), [True, np.NA, False])
+    assert_array_equal(a.all(axis=1), [False, True])
+    assert_array_equal(a.all(axis=0, skipna=True), [True, True, False])
+    assert_array_equal(a.all(axis=1, skipna=True), [False, True])
+
 def test_count_reduce_items():
     # np.count_reduce_items