MAINT: add __array_priority__ special case to masked array binary ops

njsmith · njsmith · commit c3c819ae27c1 · 2015-05-12T01:35:37.000-07:00
ndarray special methods like __add__ have a special case where if the
right argument is not an ndarray or subclass, and it has higher
__array_priority__ than the left argument, then we return
NotImplemented and let the right argument handle the operation.

ufuncs have traditionally had a similar but different special case,
where if it's a 2 input - 1 output ufunc, and the right argument is
not an ndarray (exactly, subclasses don't count), and when converted
to an ndarray ends up as an object array (presumably b/c it doesn't
have a meaningful coercion route, though who knows), and it has a
higher __array_priority__ than the left argument AND it has a
__r&lt;operation&gt;__ attribute, then they return NotImplemented.

In practice this latter special case is not used by regular ndarrays,
b/c anytime it would need to be triggered, the former special case
triggers first and the ufunc is never called. However, numpy.ma did
not have the former special case, and was thus relying on the ufunc
special case. This commit adds the special case to the numpy.ma
special methods directly, so that they no longer depend on the quirky
ufunc behaviour.

It also cleans up the relevant test to things that actually should be
true in general, instead of just testing some implementation details.
diff --git a/numpy/core/src/multiarray/number.c b/numpy/core/src/multiarray/number.c
@@ -304,6 +304,10 @@ PyArray_GenericBinaryFunction(PyArrayObject *m1, PyObject *m2, PyObject *op)
            * See also:
            * - https://github.com/numpy/numpy/issues/3502
            * - https://github.com/numpy/numpy/issues/3503
+           *
+           * NB: there's another copy of this code in
+           *    numpy.ma.core.MaskedArray._delegate_binop
+           * which should possibly be updated when this is.
            */
           double m1_prio = PyArray_GetPriority((PyObject *)m1,
                                                NPY_SCALAR_PRIORITY);
diff --git a/numpy/ma/core.py b/numpy/ma/core.py
@@ -3678,6 +3678,16 @@ def __repr__(self):
             return _print_templates['short_std'] % parameters
         return _print_templates['long_std'] % parameters
 
+    def _delegate_binop(self, other):
+        # This emulates the logic in
+        #     multiarray/number.c:PyArray_GenericBinaryFunction
+        if (not isinstance(other, np.ndarray)
+            and not hasattr(other, "__numpy_ufunc__")):
+            other_priority = getattr(other, "__array_priority__", -1000000)
+            if self.__array_priority__ < other_priority:
+                return True
+        return False
+
     def __eq__(self, other):
         "Check whether other equals self elementwise"
         if self is masked:
@@ -3746,6 +3756,8 @@ def __ne__(self, other):
     #
     def __add__(self, other):
         "Add other to self, and return a new masked array."
+        if self._delegate_binop(other):
+            return NotImplemented
         return add(self, other)
     #
     def __radd__(self, other):
@@ -3754,6 +3766,8 @@ def __radd__(self, other):
     #
     def __sub__(self, other):
         "Subtract other to self, and return a new masked array."
+        if self._delegate_binop(other):
+            return NotImplemented
         return subtract(self, other)
     #
     def __rsub__(self, other):
@@ -3762,6 +3776,8 @@ def __rsub__(self, other):
     #
     def __mul__(self, other):
         "Multiply other by self, and return a new masked array."
+        if self._delegate_binop(other):
+            return NotImplemented
         return multiply(self, other)
     #
     def __rmul__(self, other):
@@ -3770,10 +3786,14 @@ def __rmul__(self, other):
     #
     def __div__(self, other):
         "Divide other into self, and return a new masked array."
+        if self._delegate_binop(other):
+            return NotImplemented
         return divide(self, other)
     #
     def __truediv__(self, other):
         "Divide other into self, and return a new masked array."
+        if self._delegate_binop(other):
+            return NotImplemented
         return true_divide(self, other)
     #
     def __rtruediv__(self, other):
@@ -3782,6 +3802,8 @@ def __rtruediv__(self, other):
     #
     def __floordiv__(self, other):
         "Divide other into self, and return a new masked array."
+        if self._delegate_binop(other):
+            return NotImplemented
         return floor_divide(self, other)
     #
     def __rfloordiv__(self, other):
@@ -3790,6 +3812,8 @@ def __rfloordiv__(self, other):
     #
     def __pow__(self, other):
         "Raise self to the power other, masking the potential NaNs/Infs"
+        if self._delegate_binop(other):
+            return NotImplemented
         return power(self, other)
     #
     def __rpow__(self, other):
diff --git a/numpy/ma/tests/test_core.py b/numpy/ma/tests/test_core.py
@@ -12,6 +12,7 @@
 import sys
 import pickle
 from functools import reduce
+import operator
 
 from nose.tools import assert_raises
 
@@ -1691,16 +1692,15 @@ def test_ndarray_mask(self):
         self.assertTrue(not isinstance(test.mask, MaskedArray))
 
     def test_treatment_of_NotImplemented(self):
-        # Check any NotImplemented returned by umath.<ufunc> is passed on
+        # Check that NotImplemented is returned at appropriate places
+
         a = masked_array([1., 2.], mask=[1, 0])
-        # basic tests for _MaskedBinaryOperation
-        assert_(a.__mul__('abc') is NotImplemented)
-        assert_(multiply.outer(a, 'abc') is NotImplemented)
-        # and for _DomainedBinaryOperation
-        assert_(a.__div__('abc') is NotImplemented)
-
-        # also check explicitly that rmul of another class can be accessed
-        class MyClass(str):
+        self.assertRaises(TypeError, operator.mul, a, "abc")
+        self.assertRaises(TypeError, operator.truediv, a, "abc")
+
+        class MyClass(object):
+            __array_priority__ = a.__array_priority__ + 1
+
             def __mul__(self, other):
                 return "My mul"
 
