|
"""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__`. |
|
|
|
""" |
|
|
|
|
|
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 |
|
|
