"""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