codekingpro/portable-devtools

1"""Tools for testing implementations of __array_function__ and ufunc overrides
2
3
4"""
5
6import numpy._core.umath as _umath
7from numpy import ufunc as _ufunc
8from numpy._core.overrides import ARRAY_FUNCTIONS as _array_functions
9
10
11def get_overridable_numpy_ufuncs():
12    """List all numpy ufuncs overridable via `__array_ufunc__`
13
14    Parameters
15    ----------
16    None
17
18    Returns
19    -------
20    set
21        A set containing all overridable ufuncs in the public numpy API.
22    """
23    ufuncs = {obj for obj in _umath.__dict__.values()
24              if isinstance(obj, _ufunc)}
25    return ufuncs
26
27
28def allows_array_ufunc_override(func):
29    """Determine if a function can be overridden via `__array_ufunc__`
30
31    Parameters
32    ----------
33    func : callable
34        Function that may be overridable via `__array_ufunc__`
35
36    Returns
37    -------
38    bool
39        `True` if `func` is overridable via `__array_ufunc__` and
40        `False` otherwise.
41
42    Notes
43    -----
44    This function is equivalent to ``isinstance(func, np.ufunc)`` and
45    will work correctly for ufuncs defined outside of Numpy.
46
47    """
48    return isinstance(func, _ufunc)
49
50
51def get_overridable_numpy_array_functions():
52    """List all numpy functions overridable via `__array_function__`
53
54    Parameters
55    ----------
56    None
57
58    Returns
59    -------
60    set
61        A set containing all functions in the public numpy API that are
62        overridable via `__array_function__`.
63
64    """
65    # 'import numpy' doesn't import recfunctions, so make sure it's imported
66    # so ufuncs defined there show up in the ufunc listing
67    from numpy.lib import recfunctions  # noqa: F401
68    return _array_functions.copy()
69
70def allows_array_function_override(func):
71    """Determine if a Numpy function can be overridden via `__array_function__`
72
73    Parameters
74    ----------
75    func : callable
76        Function that may be overridable via `__array_function__`
77
78    Returns
79    -------
80    bool
81        `True` if `func` is a function in the Numpy API that is
82        overridable via `__array_function__` and `False` otherwise.
83    """
84    return func in _array_functions
85