"""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)} return ufuncs def allows_array_ufunc_override(func): """Determine if a function can be overridden 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. Notes ----- 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 overridden 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