API: Add numpy.testing.overrides to aid testing of custom array containers

Closes #15544

3 files changed, 87 insertions, 0 deletions

diff --git a/numpy/core/overrides.py b/numpy/core/overrides.py
index 450464f89..6d3680915 100644
--- a/numpy/core/overrides.py
+++ b/numpy/core/overrides.py
@@ -8,6 +8,8 @@ from numpy.core._multiarray_umath import (
 from numpy.compat._inspect import getargspec
 
 
+ARRAY_FUNCTIONS = set()
+
 ARRAY_FUNCTION_ENABLED = bool(
     int(os.environ.get('NUMPY_EXPERIMENTAL_ARRAY_FUNCTION', 1)))
 
@@ -208,6 +210,8 @@ def array_function_dispatch(dispatcher, module=None, verify=True,
 
         public_api._implementation = implementation
 
+        ARRAY_FUNCTIONS.add(public_api)
+
         return public_api
 
     return decorator
diff --git a/numpy/testing/overrides.py b/numpy/testing/overrides.py
new file mode 100644
index 000000000..d20ed60e5
--- /dev/null
+++ b/numpy/testing/overrides.py
@@ -0,0 +1,82 @@
+"""Tools for testing implementations of __array_function__ and ufunc overrides
+
+
+"""
+
+from numpy.core.overrides import ARRAY_FUNCTIONS as _array_functions
+from numpy import ufunc as _ufunc
+import numpy.core.umath as _umath
+
+def get_overridable_numpy_ufuncs():
+    """List all numpy ufuncs overridable via `__array_ufunc__`
+
+    Parameters
+    ----------
+    None
+
+    Returns
+    -------
+    set
+        A set containing all overridable ufuncs in the public numpy API.
+    """
+    ufuncs = {obj for obj in _umath.__dict__.values()
+              if isinstance(obj, _ufunc)}
+    
+
+def allows_array_ufunc_override(func):
+    """Determine if a function can be overriden via `__array_ufunc__`
+
+    Parameters
+    ----------
+    func : callable
+        Function that may be overridable via `__array_ufunc__`
+
+    Returns
+    -------
+    bool
+        `True` if `func` is overridable via `__array_ufunc__` and
+        `False` otherwise.
+
+    Note
+    ----
+    This function is equivalent to `isinstance(func, np.ufunc)` and
+    will work correctly for ufuncs defined outside of Numpy.
+
+    """
+    return isinstance(func, np.ufunc)
+
+
+def get_overridable_numpy_array_functions():
+    """List all numpy functions overridable via `__array_function__`
+
+    Parameters
+    ----------
+    None
+
+    Returns
+    -------
+    set
+        A set containing all functions in the public numpy API that are
+        overridable via `__array_function__`.
+
+    """
+    # 'import numpy' doesn't import recfunctions, so make sure it's imported
+    # so ufuncs defined there show up in the ufunc listing
+    from numpy.lib import recfunctions
+    return _array_functions.copy()
+
+def allows_array_function_override(func):
+    """Determine if a Numpy function can be overriden via `__array_function__`
+
+    Parameters
+    ----------
+    func : callable
+        Function that may be overridable via `__array_function__`
+
+    Returns
+    -------
+    bool
+        `True` if `func` is a function in the Numpy API that is
+        overridable via `__array_function__` and `False` otherwise.
+    """
+    return func in _array_functions
diff --git a/numpy/tests/test_public_api.py b/numpy/tests/test_public_api.py
index 92a34bf91..d04e24218 100644
--- a/numpy/tests/test_public_api.py
+++ b/numpy/tests/test_public_api.py
@@ -157,6 +157,7 @@ PUBLIC_MODULES = ['numpy.' + s for s in [
     "polynomial.polynomial",
     "random",
     "testing",
+    "testing.overrides",
     "typing",
     "typing.mypy_plugin",
     "version",