Merge pull request #3907 from mhvk/ma/allow-subclass-in-ufunc

BUG allow subclasses in MaskedArray ufuncs -- for non-ndarray _data

3 files changed, 147 insertions, 104 deletions

diff --git a/numpy/ma/core.py b/numpy/ma/core.py
index 8e14d7c1a..24e18a9d6 100644
--- a/numpy/ma/core.py
+++ b/numpy/ma/core.py
@@ -880,14 +880,10 @@ class _MaskedUnaryOperation:
             except TypeError:
                 pass
         # Transform to
-        if isinstance(a, MaskedArray):
-            subtype = type(a)
-        else:
-            subtype = MaskedArray
-        result = result.view(subtype)
-        result._mask = m
-        result._update_from(a)
-        return result
+        masked_result = result.view(get_masked_subclass(a))
+        masked_result._mask = m
+        masked_result._update_from(result)
+        return masked_result
     #
     def __str__ (self):
         return "Masked version of %s. [Invalid values are masked]" % str(self.f)
@@ -928,8 +924,12 @@ class _MaskedBinaryOperation:
     def __call__ (self, a, b, *args, **kwargs):
         "Execute the call behavior."
         # Get the data, as ndarray
-        (da, db) = (getdata(a, subok=False), getdata(b, subok=False))
-        # Get the mask
+        (da, db) = (getdata(a), getdata(b))
+        # Get the result
+        with np.errstate():
+            np.seterr(divide='ignore', invalid='ignore')
+            result = self.f(da, db, *args, **kwargs)
+        # Get the mask for the result
         (ma, mb) = (getmask(a), getmask(b))
         if ma is nomask:
             if mb is nomask:
@@ -940,12 +940,6 @@ class _MaskedBinaryOperation:
             m = umath.logical_or(ma, getmaskarray(b))
         else:
             m = umath.logical_or(ma, mb)
-        # Get the result
-        with np.errstate(divide='ignore', invalid='ignore'):
-            result = self.f(da, db, *args, **kwargs)
-        # check it worked
-        if result is NotImplemented:
-            return NotImplemented
         # Case 1. : scalar
         if not result.ndim:
             if m:
@@ -953,28 +947,26 @@ class _MaskedBinaryOperation:
             return result
         # Case 2. : array
         # Revert result to da where masked
-        if m is not nomask:
-            np.copyto(result, da, casting='unsafe', where=m)
+        if m is not nomask and m.any():
+            # any errors, just abort; impossible to guarantee masked values
+            try:
+                np.copyto(result, 0, casting='unsafe', where=m)
+                # avoid using "*" since this may be overlaid
+                masked_da = umath.multiply(m, da)
+                # only add back if it can be cast safely
+                if np.can_cast(masked_da.dtype, result.dtype, casting='safe'):
+                    result += masked_da
+            except:
+                pass
         # Transforms to a (subclass of) MaskedArray
-        result = result.view(get_masked_subclass(a, b))
-        result._mask = m
-        # Update the optional info from the inputs
-        if isinstance(b, MaskedArray):
-            if isinstance(a, MaskedArray):
-                result._update_from(a)
-            else:
-                result._update_from(b)
-        elif isinstance(a, MaskedArray):
-            result._update_from(a)
-        return result
-
+        masked_result = result.view(get_masked_subclass(a, b))
+        masked_result._mask = m
+        masked_result._update_from(result)
+        return masked_result
 
     def reduce(self, target, axis=0, dtype=None):
         """Reduce `target` along the given `axis`."""
-        if isinstance(target, MaskedArray):
-            tclass = type(target)
-        else:
-            tclass = MaskedArray
+        tclass = get_masked_subclass(target)
         m = getmask(target)
         t = filled(target, self.filly)
         if t.shape == ():
@@ -982,24 +974,30 @@ class _MaskedBinaryOperation:
             if m is not nomask:
                 m = make_mask(m, copy=1)
                 m.shape = (1,)
+
         if m is nomask:
-            return self.f.reduce(t, axis).view(tclass)
-        t = t.view(tclass)
-        t._mask = m
-        tr = self.f.reduce(getdata(t), axis, dtype=dtype or t.dtype)
-        mr = umath.logical_and.reduce(m, axis)
-        tr = tr.view(tclass)
-        if mr.ndim > 0:
-            tr._mask = mr
-            return tr
-        elif mr:
-            return masked
-        return tr
+            tr = self.f.reduce(t, axis)
+            mr = nomask
+        else:
+            tr = self.f.reduce(t, axis, dtype=dtype or t.dtype)
+            mr = umath.logical_and.reduce(m, axis)
 
-    def outer (self, a, b):
+        if not tr.shape:
+            if mr:
+                return masked
+            else:
+                return tr
+        masked_tr = tr.view(tclass)
+        masked_tr._mask = mr
+        masked_tr._update_from(tr)
+        return masked_tr
+
+    def outer(self, a, b):
         """Return the function applied to the outer product of a and b.
 
         """
+        (da, db) = (getdata(a), getdata(b))
+        d = self.f.outer(da, db)
         ma = getmask(a)
         mb = getmask(b)
         if ma is nomask and mb is nomask:
@@ -1010,31 +1008,28 @@ class _MaskedBinaryOperation:
             m = umath.logical_or.outer(ma, mb)
         if (not m.ndim) and m:
             return masked
-        (da, db) = (getdata(a), getdata(b))
-        d = self.f.outer(da, db)
-        # check it worked
-        if d is NotImplemented:
-            return NotImplemented
         if m is not nomask:
             np.copyto(d, da, where=m)
-        if d.shape:
-            d = d.view(get_masked_subclass(a, b))
-            d._mask = m
-        return d
+        if not d.shape:
+            return d
+        masked_d = d.view(get_masked_subclass(a, b))
+        masked_d._mask = m
+        masked_d._update_from(d)
+        return masked_d
 
-    def accumulate (self, target, axis=0):
+    def accumulate(self, target, axis=0):
         """Accumulate `target` along `axis` after filling with y fill
         value.
 
         """
-        if isinstance(target, MaskedArray):
-            tclass = type(target)
-        else:
-            tclass = MaskedArray
+        tclass = get_masked_subclass(target)
         t = filled(target, self.filly)
-        return self.f.accumulate(t, axis).view(tclass)
+        result = self.f.accumulate(t, axis)
+        masked_result = result.view(tclass)
+        masked_result._update_from(result)
+        return masked_result
 
-    def __str__ (self):
+    def __str__(self):
         return "Masked version of " + str(self.f)
 
 
@@ -1074,19 +1069,15 @@ class _DomainedBinaryOperation:
 
     def __call__(self, a, b, *args, **kwargs):
         "Execute the call behavior."
-        # Get the data and the mask
-        (da, db) = (getdata(a, subok=False), getdata(b, subok=False))
-        (ma, mb) = (getmask(a), getmask(b))
+        # Get the data
+        (da, db) = (getdata(a), getdata(b))
         # Get the result
         with np.errstate(divide='ignore', invalid='ignore'):
             result = self.f(da, db, *args, **kwargs)
-        # check it worked
-        if result is NotImplemented:
-            return NotImplemented
-        # Get the mask as a combination of ma, mb and invalid
+        # Get the mask as a combination of the source masks and invalid
         m = ~umath.isfinite(result)
-        m |= ma
-        m |= mb
+        m |= getmask(a)
+        m |= getmask(b)
         # Apply the domain
         domain = ufunc_domain.get(self.f, None)
         if domain is not None:
@@ -1097,18 +1088,23 @@ class _DomainedBinaryOperation:
                 return masked
             else:
                 return result
-        # When the mask is True, put back da
-        np.copyto(result, da, casting='unsafe', where=m)
-        result = result.view(get_masked_subclass(a, b))
-        result._mask = m
-        if isinstance(b, MaskedArray):
-            if isinstance(a, MaskedArray):
-                result._update_from(a)
-            else:
-                result._update_from(b)
-        elif isinstance(a, MaskedArray):
-            result._update_from(a)
-        return result
+        # When the mask is True, put back da if possible
+        # any errors, just abort; impossible to guarantee masked values
+        try:
+            np.copyto(result, 0, casting='unsafe', where=m)
+            # avoid using "*" since this may be overlaid
+            masked_da = umath.multiply(m, da)
+            # only add back if it can be cast safely
+            if np.can_cast(masked_da.dtype, result.dtype, casting='safe'):
+                result += masked_da
+        except:
+            pass
+
+        # Transforms to a (subclass of) MaskedArray
+        masked_result = result.view(get_masked_subclass(a, b))
+        masked_result._mask = m
+        masked_result._update_from(result)
+        return masked_result
 
     def __str__ (self):
         return "Masked version of " + str(self.f)
@@ -1361,7 +1357,7 @@ def getmaskarray(arr):
     """
     mask = getmask(arr)
     if mask is nomask:
-        mask = make_mask_none(np.shape(arr), getdata(arr).dtype)
+        mask = make_mask_none(np.shape(arr), getattr(arr, 'dtype', None))
     return mask
 
 def is_mask(m):
@@ -3756,34 +3752,38 @@ class MaskedArray(ndarray):
         return check
     #
     def __add__(self, other):
-        "Add other to self, and return a new masked array."
+        "Add self to other, and return a new masked array."
         if self._delegate_binop(other):
             return NotImplemented
         return add(self, other)
     #
     def __radd__(self, other):
         "Add other to self, and return a new masked array."
-        return add(self, other)
+        # In analogy with __rsub__ and __rdiv__, use original order:
+        # we get here from `other + self`.
+        return add(other, self)
     #
     def __sub__(self, other):
-        "Subtract other to self, and return a new masked array."
+        "Subtract other from self, and return a new masked array."
         if self._delegate_binop(other):
             return NotImplemented
         return subtract(self, other)
     #
     def __rsub__(self, other):
-        "Subtract other to self, and return a new masked array."
+        "Subtract self from other, and return a new masked array."
         return subtract(other, self)
     #
     def __mul__(self, other):
-        "Multiply other by self, and return a new masked array."
+        "Multiply self by other, and return a new masked array."
         if self._delegate_binop(other):
             return NotImplemented
         return multiply(self, other)
     #
     def __rmul__(self, other):
         "Multiply other by self, and return a new masked array."
-        return multiply(self, other)
+        # In analogy with __rsub__ and __rdiv__, use original order:
+        # we get here from `other * self`.
+        return multiply(other, self)
     #
     def __div__(self, other):
         "Divide other into self, and return a new masked array."
@@ -3798,7 +3798,7 @@ class MaskedArray(ndarray):
         return true_divide(self, other)
     #
     def __rtruediv__(self, other):
-        "Divide other into self, and return a new masked array."
+        "Divide self into other, and return a new masked array."
         return true_divide(other, self)
     #
     def __floordiv__(self, other):
@@ -3808,7 +3808,7 @@ class MaskedArray(ndarray):
         return floor_divide(self, other)
     #
     def __rfloordiv__(self, other):
-        "Divide other into self, and return a new masked array."
+        "Divide self into other, and return a new masked array."
         return floor_divide(other, self)
     #
     def __pow__(self, other):
@@ -3818,7 +3818,7 @@ class MaskedArray(ndarray):
         return power(self, other)
     #
     def __rpow__(self, other):
-        "Raise self to the power other, masking the potential NaNs/Infs"
+        "Raise other to the power self, masking the potential NaNs/Infs"
         return power(other, self)
     #............................................
     def __iadd__(self, other):
diff --git a/numpy/ma/tests/test_core.py b/numpy/ma/tests/test_core.py
index 909c27c90..9cc784d02 100644
--- a/numpy/ma/tests/test_core.py
+++ b/numpy/ma/tests/test_core.py
@@ -1754,6 +1754,29 @@ class TestUfuncs(TestCase):
         assert_(me * a == "My mul")
         assert_(a * me == "My rmul")
 
+        # and that __array_priority__ is respected
+        class MyClass2(object):
+            __array_priority__ = 100
+
+            def __mul__(self, other):
+                return "Me2mul"
+
+            def __rmul__(self, other):
+                return "Me2rmul"
+
+            def __rdiv__(self, other):
+                return "Me2rdiv"
+
+            __rtruediv__ = __rdiv__
+
+        me_too = MyClass2()
+        assert_(a.__mul__(me_too) is NotImplemented)
+        assert_(all(multiply.outer(a, me_too) == "Me2rmul"))
+        assert_(a.__truediv__(me_too) is NotImplemented)
+        assert_(me_too * a == "Me2mul")
+        assert_(a * me_too == "Me2rmul")
+        assert_(a / me_too == "Me2rdiv")
+
 
 #------------------------------------------------------------------------------
 class TestMaskedArrayInPlaceArithmetics(TestCase):
diff --git a/numpy/ma/tests/test_subclassing.py b/numpy/ma/tests/test_subclassing.py
index 07fc8fdd6..2e98348e6 100644
--- a/numpy/ma/tests/test_subclassing.py
+++ b/numpy/ma/tests/test_subclassing.py
@@ -24,18 +24,27 @@ class SubArray(np.ndarray):
     # in the  dictionary `info`.
     def __new__(cls,arr,info={}):
         x = np.asanyarray(arr).view(cls)
-        x.info = info
+        x.info = info.copy()
         return x
 
     def __array_finalize__(self, obj):
-        self.info = getattr(obj, 'info', {})
+        if callable(getattr(super(SubArray, self),
+                            '__array_finalize__', None)):
+            super(SubArray, self).__array_finalize__(obj)
+        self.info = getattr(obj, 'info', {}).copy()
         return
 
     def __add__(self, other):
-        result = np.ndarray.__add__(self, other)
-        result.info.update({'added':result.info.pop('added', 0)+1})
+        result = super(SubArray, self).__add__(other)
+        result.info['added'] = result.info.get('added', 0) + 1
         return result
 
+    def __iadd__(self, other):
+        result = super(SubArray, self).__iadd__(other)
+        result.info['iadded'] = result.info.get('iadded', 0) + 1
+        return result
+
+
 subarray = SubArray
 
 
@@ -47,11 +56,6 @@ class MSubArray(SubArray, MaskedArray):
         _data.info = subarr.info
         return _data
 
-    def __array_finalize__(self, obj):
-        MaskedArray.__array_finalize__(self, obj)
-        SubArray.__array_finalize__(self, obj)
-        return
-
     def _get_series(self):
         _view = self.view(MaskedArray)
         _view._sharedmask = False
@@ -82,8 +86,11 @@ class MMatrix(MaskedArray, np.matrix,):
 mmatrix = MMatrix
 
 
-# also a subclass that overrides __str__, __repr__ and __setitem__, disallowing
+# Also a subclass that overrides __str__, __repr__ and __setitem__, disallowing
 # setting to non-class values (and thus np.ma.core.masked_print_option)
+# and overrides __array_wrap__, updating the info dict, to check that this
+# doesn't get destroyed by MaskedArray._update_from.  But this one also needs
+# its own iterator...
 class CSAIterator(object):
     """
     Flat iterator object that uses its own setter/getter
@@ -150,6 +157,13 @@ class ComplicatedSubArray(SubArray):
         y = self.ravel()
         y[:] = value
 
+    def __array_wrap__(self, obj, context=None):
+        obj = super(ComplicatedSubArray, self).__array_wrap__(obj, context)
+        if context is not None and context[0] is np.multiply:
+            obj.info['multiplied'] = obj.info.get('multiplied', 0) + 1
+
+        return obj
+
 
 class TestSubclassing(TestCase):
     # Test suite for masked subclasses of ndarray.
@@ -218,6 +232,12 @@ class TestSubclassing(TestCase):
         self.assertTrue(isinstance(z, MSubArray))
         self.assertTrue(isinstance(z._data, SubArray))
         self.assertTrue(z._data.info['added'] > 0)
+        # Test that inplace methods from data get used (gh-4617)
+        ym += 1
+        self.assertTrue(isinstance(ym, MaskedArray))
+        self.assertTrue(isinstance(ym, MSubArray))
+        self.assertTrue(isinstance(ym._data, SubArray))
+        self.assertTrue(ym._data.info['iadded'] > 0)
         #
         ym._set_mask([1, 0, 0, 0, 1])
         assert_equal(ym._mask, [1, 0, 0, 0, 1])