Merge pull request #8508 from eric-wieser/fix-masked-ufunc-shape

BUG: Fix loss of dimensionality of np.ma.masked in ufunc

2 files changed, 27 insertions, 2 deletions

diff --git a/numpy/ma/core.py b/numpy/ma/core.py
index 4466dc0af..49c790d4f 100644
--- a/numpy/ma/core.py
+++ b/numpy/ma/core.py
@@ -6130,8 +6130,11 @@ class MaskedConstant(MaskedArray):
     def __array_finalize__(self, obj):
         return
 
-    def __array_wrap__(self, obj):
-        return self
+    def __array_prepare__(self, obj, context=None):
+        return self.view(MaskedArray).__array_prepare__(obj, context)
+
+    def __array_wrap__(self, obj, context=None):
+        return self.view(MaskedArray).__array_wrap__(obj, context)
 
     def __str__(self):
         return str(masked_print_option._display)
diff --git a/numpy/ma/tests/test_core.py b/numpy/ma/tests/test_core.py
index b5be24b31..c6f276211 100644
--- a/numpy/ma/tests/test_core.py
+++ b/numpy/ma/tests/test_core.py
@@ -4425,6 +4425,28 @@ class TestOptionalArgs(TestCase):
         assert_raises(ValueError, count, np.ma.array(1), axis=1)
 
 
+class TestMaskedConstant(TestCase):
+    def _do_add_test(self, add):
+        # sanity check
+        self.assertIs(add(np.ma.masked, 1), np.ma.masked)
+
+        # now try with a vector
+        vector = np.array([1, 2, 3])
+        result = add(np.ma.masked, vector)
+
+        # lots of things could go wrong here
+        assert_(result is not np.ma.masked)
+        assert_(not isinstance(result, np.ma.core.MaskedConstant))
+        assert_equal(result.shape, vector.shape)
+        assert_equal(np.ma.getmask(result), np.ones(vector.shape, dtype=bool))
+
+    def test_ufunc(self):
+        self._do_add_test(np.add)
+
+    def test_operator(self):
+        self._do_add_test(lambda a, b: a + b)
+
+
 def test_masked_array():
     a = np.ma.array([0, 1, 2, 3], mask=[0, 0, 1, 0])
     assert_equal(np.argwhere(a), [[1], [3]])