MAINT: fixes and tests for assert_array_equal on subclasses

numpy · shoyer · Oct 23, 2018 · Oct 8, 2018 · Oct 9, 2018 · Oct 12, 2018
commit 3844ade67d55a69ee64afd810a95bc522ec74b7e
diff --git a/numpy/testing/_private/utils.py b/numpy/testing/_private/utils.py
@@ -707,23 +707,28 @@ def func_assert_same_pos(x, y, func=isnan, hasval='nan'):
         """
         x_id = func(x)
         y_id = func(y)
-        # Cast to bool_ to ensure than an all() method is available, even for
-        # subclasses that define equality to return Python booleans. We are not
-        # committed to supporting such subclasses, but they used to work.
-        # We don't use np.all() because it's more likely that that method is
-        # not supported on subclasses which implement `__array_function__`.
-        if bool_(x_id == y_id).all():
+        # We include work-arounds here to handle three types of slightly
+        # pathological ndarray subclasses:
+        # (1) all() on `masked` array scalars can return masked arrays, so we
+        #     use != True
+        # (2) __eq__ on some ndarray subclasses returns Python booleans
+        #     instead of element-wise comparisons, so we cast to bool_() and
+        #     use isinstance(..., bool) checks
+        # (3) subclasses with bare-bones __array_function__ implemenations may
+        #     not implement np.all(), so favor using the .all() method
+        # We are not committed to supporting such subclasses, but it's nice to
+        # support them if possible.
+        if bool_(x_id == y_id).all() != True:
             msg = build_err_msg([x, y],
                                 err_msg + '\nx and y %s location mismatch:'
                                 % (hasval), verbose=verbose, header=header,
                                 names=('x', 'y'), precision=precision)
             raise AssertionError(msg)
         # If there is a scalar, then here we know the array has the same
         # flag as it everywhere, so we should return the scalar flag.
-        # Cast to bool_ to avoid returning `masked` arrays.
-        if x_id.ndim == 0:
+        if isinstance(x_id, bool) or x_id.ndim == 0:
             return bool_(x_id)
-        elif y_id.ndim == 0:
+        elif isinstance(x_id, bool) or y_id.ndim == 0:
             return bool_(y_id)
         else:
             return y_id

diff --git a/numpy/testing/tests/test_utils.py b/numpy/testing/tests/test_utils.py
@@ -158,6 +158,44 @@ def test_masked_nan_inf(self):
         self._test_equal(a, b)
         self._test_equal(b, a)
 
+    def test_subclass_that_overrides_eq(self):
+        # While we cannot guarantee testing functions will always work for
+        # subclasses, the tests should ideally rely only on subclasses having
+        # comparison operators, not on them being able to store booleans
+        # (which, e.g., astropy Quantity cannot usefully do). See gh-8452.
+        class MyArray(np.ndarray):
+            def __eq__(self, other):
+                return bool(np.equal(self, other).all())
+
+            def __ne__(self, other):
+                return not self == other
+
+        a = np.array([1., 2.]).view(MyArray)
+        b = np.array([2., 3.]).view(MyArray)
+        assert_(type(a == a), bool)
+        assert_(a == a)
+        assert_(a != b)
+        self._test_equal(a, a)
+        self._test_not_equal(a, b)
+        self._test_not_equal(b, a)
+
+    def test_subclass_that_does_not_implement_npall(self):
+        # While we cannot guarantee testing functions will always work for
+        # subclasses, the tests should ideally rely only on subclasses having
+        # comparison operators, not on them being able to store booleans
+        # (which, e.g., astropy Quantity cannot usefully do). See gh-8452.
+        class MyArray(np.ndarray):
+            def __array_function__(self, *args, **kwargs):
+                return NotImplemented
+
+        a = np.array([1., 2.]).view(MyArray)
+        b = np.array([2., 3.]).view(MyArray)
+        with assert_raises(TypeError):
+            np.all(a)
+        self._test_equal(a, a)
+        self._test_not_equal(a, b)
+        self._test_not_equal(b, a)
+
 
 class TestBuildErrorMessage(object):