BUG: fix np.ma.masked_where(copy=False) when input has no mask (#18967)

constanzafierro · web-flow · commit b6eb3d8e64e8 · 2021-05-10T07:25:02.000+02:00
Fixes gh-18946
diff --git a/numpy/ma/core.py b/numpy/ma/core.py
@@ -1936,6 +1936,10 @@ def masked_where(condition, a, copy=True):
     result = a.view(cls)
     # Assign to *.mask so that structured masks are handled correctly.
     result.mask = _shrink_mask(cond)
+    # There is no view of a boolean so when 'a' is a MaskedArray with nomask
+    # the update to the result's mask has no effect.
+    if not copy and hasattr(a, '_mask') and getmask(a) is nomask:
+        a._mask = result._mask.view()
     return result
 
 
diff --git a/numpy/ma/tests/test_core.py b/numpy/ma/tests/test_core.py
@@ -5175,6 +5175,16 @@ def test_masked_array():
     a = np.ma.array([0, 1, 2, 3], mask=[0, 0, 1, 0])
     assert_equal(np.argwhere(a), [[1], [3]])
 
+def test_masked_array_no_copy():
+    # check nomask array is updated in place
+    a = np.ma.array([1, 2, 3, 4])
+    _ = np.ma.masked_where(a == 3, a, copy=False)
+    assert_array_equal(a.mask, [False, False, True, False])
+    # check masked array is updated in place
+    a = np.ma.array([1, 2, 3, 4], mask=[1, 0, 0, 0])
+    _ = np.ma.masked_where(a == 3, a, copy=False)
+    assert_array_equal(a.mask, [True, False, True, False])
+
 def test_append_masked_array():
     a = np.ma.masked_equal([1,2,3], value=2)
     b = np.ma.masked_equal([4,3,2], value=2)
@@ -5213,7 +5223,6 @@ def test_append_masked_array_along_axis():
     assert_array_equal(result.data, expected.data)
     assert_array_equal(result.mask, expected.mask)
 
-
 def test_default_fill_value_complex():
     # regression test for Python 3, where 'unicode' was not defined
     assert_(default_fill_value(1 + 1j) == 1.e20 + 0.0j)
