File size: 2,123 Bytes
fe41391 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
"""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
|