Fix precedence of ndarray subclasses and misc cleanup

numpy · charris · Sep 24, 2018 · Sep 21, 2018 · Sep 21, 2018 · Sep 21, 2018
commit 0da1b95ea9180ab613eccc80ebe39eb4f48d99d3
diff --git a/numpy/core/overrides.py b/numpy/core/overrides.py
@@ -6,6 +6,9 @@
 from numpy.core.multiarray import ndarray
 
 
+_NDARRAY_ARRAY_FUNCTION = ndarray.__array_function__
+
+
 def get_overloaded_types_and_args(relevant_args):
     """Returns a list of arguments on which to call __array_function__.
 
@@ -17,36 +20,32 @@ def get_overloaded_types_and_args(relevant_args):
     overloaded_args = []
     for arg in relevant_args:
         arg_type = type(arg)
-        if arg_type not in overloaded_types:
-            try:
-                array_function = arg_type.__array_function__
-            except AttributeError:
-                continue
+        if (arg_type not in overloaded_types and
+                hasattr(arg_type, '__array_function__')):
 
             overloaded_types.append(arg_type)
 
-            if array_function is not ndarray.__array_function__:
-                # By default, insert this argument at the end, but if it is
-                # subclass of another argument, insert it before that argument.
-                # This ensures "subclasses before superclasses".
-                index = len(overloaded_args)
-                for i, old_arg in enumerate(overloaded_args):
-                    if issubclass(arg_type, type(old_arg)):
-                        index = i
-                        break
-                overloaded_args.insert(index, arg)
+            # By default, insert this argument at the end, but if it is
+            # subclass of another argument, insert it before that argument.
+            # This ensures "subclasses before superclasses".
+            index = len(overloaded_args)
+            for i, old_arg in enumerate(overloaded_args):
+                if issubclass(arg_type, type(old_arg)):
+                    index = i
+                    break
+            overloaded_args.insert(index, arg)
 
-    return tuple(overloaded_types), tuple(overloaded_args)
+    # Special handling for ndarray.
+    overloaded_args = [
+        arg for arg in overloaded_args
+        if type(arg).__array_function__ is not _NDARRAY_ARRAY_FUNCTION
+    ]
 
+    return overloaded_types, overloaded_args
 
-def try_array_function_override(func, relevant_arguments, args, kwargs):
-    # TODO: consider simplifying the interface, to only require either `types`
-    # (by calling __array_function__ a classmethod) or `overloaded_args` (by
-    # dropping `types` from the signature of __array_function__)
-    types, overloaded_args = get_overloaded_types_and_args(relevant_arguments)
-    if not overloaded_args:
-        return False, None
 
+def array_function_override(overloaded_args, func, types, args, kwargs):
+    """Call __array_function__ implementations."""
     for overloaded_arg in overloaded_args:
         # Note that we're only calling __array_function__ on the *first*
         # occurence of each argument type. This is necessary for reasonable
@@ -56,7 +55,7 @@ def try_array_function_override(func, relevant_arguments, args, kwargs):
         result = overloaded_arg.__array_function__(func, types, args, kwargs)
 
         if result is not NotImplemented:
-            return True, result
+            return result
 
     raise TypeError('no implementation found for {} on types that implement '
                     '__array_function__: {}'
@@ -68,11 +67,20 @@ def array_function_dispatch(dispatcher):
     def decorator(func):
         @functools.wraps(func)
         def new_func(*args, **kwargs):
+            # Collect array-like arguments.
             relevant_arguments = dispatcher(*args, **kwargs)
-            success, value = try_array_function_override(
-                new_func, relevant_arguments, args, kwargs)
-            if success:
-                return value
-            return func(*args, **kwargs)
+            # Check for __array_function__ methods.
+            types, overloaded_args = get_overloaded_types_and_args(
+                relevant_arguments)
+            # Call overrides, if necessary.
+            if overloaded_args:
+                # new_func is the function exposed in NumPy's public API. We
+                # use it instead of func so __array_function__ implementations
+                # can do equality/identity comparisons.
+                return array_function_override(
+                    overloaded_args, new_func, types, args, kwargs)
+            else:
+                return func(*args, **kwargs)
+
         return new_func
     return decorator
diff --git a/numpy/core/tests/test_overrides.py b/numpy/core/tests/test_overrides.py
@@ -82,6 +82,23 @@ class Other(object):
         assert_equal(set(types), {np.ndarray, Other})
         assert_equal(list(args), [other])
 
+    def test_ndarray_subclass_and_duck_array(self):
+
+        class OverrideSub(np.ndarray):
+            __array_function__ = _return_self
+
+        class Other(object):
+            __array_function__ = _return_self
+
+        array = np.array(1)
+        subarray = np.array(1).view(OverrideSub)
+        other = Other()
+
+        assert_equal(_get_overloaded_args([array, subarray, other]),
+                     [subarray, other])
+        assert_equal(_get_overloaded_args([array, other, subarray]),
+                     [subarray, other])
+
     def test_many_duck_arrays(self):
 
         class A(object):