codekingpro/portable-devtools

1"""Implementation of __array_function__ overrides from NEP-18."""
2import collections
3import functools
4import inspect
5
6from numpy._core._multiarray_umath import (
7    _ArrayFunctionDispatcher,
8    _get_implementing_args,
9    add_docstring,
10)
11from numpy._utils import set_module  # noqa: F401
12from numpy._utils._inspect import getargspec
13
14ARRAY_FUNCTIONS = set()
15
16array_function_like_doc = (
17    """like : array_like, optional
18        Reference object to allow the creation of arrays which are not
19        NumPy arrays. If an array-like passed in as ``like`` supports
20        the ``__array_function__`` protocol, the result will be defined
21        by it. In this case, it ensures the creation of an array object
22        compatible with that passed in via this argument."""
23)
24
25def get_array_function_like_doc(public_api, docstring_template=""):
26    ARRAY_FUNCTIONS.add(public_api)
27    docstring = public_api.__doc__ or docstring_template
28    return docstring.replace("${ARRAY_FUNCTION_LIKE}", array_function_like_doc)
29
30def finalize_array_function_like(public_api):
31    public_api.__doc__ = get_array_function_like_doc(public_api)
32    return public_api
33
34
35add_docstring(
36    _ArrayFunctionDispatcher,
37    """
38    Class to wrap functions with checks for __array_function__ overrides.
39
40    All arguments are required, and can only be passed by position.
41
42    Parameters
43    ----------
44    dispatcher : function or None
45        The dispatcher function that returns a single sequence-like object
46        of all arguments relevant.  It must have the same signature (except
47        the default values) as the actual implementation.
48        If ``None``, this is a ``like=`` dispatcher and the
49        ``_ArrayFunctionDispatcher`` must be called with ``like`` as the
50        first (additional and positional) argument.
51    implementation : function
52        Function that implements the operation on NumPy arrays without
53        overrides.  Arguments passed calling the ``_ArrayFunctionDispatcher``
54        will be forwarded to this (and the ``dispatcher``) as if using
55        ``*args, **kwargs``.
56
57    Attributes
58    ----------
59    _implementation : function
60        The original implementation passed in.
61    """)
62
63
64# exposed for testing purposes; used internally by _ArrayFunctionDispatcher
65add_docstring(
66    _get_implementing_args,
67    """
68    Collect arguments on which to call __array_function__.
69
70    Parameters
71    ----------
72    relevant_args : iterable of array-like
73        Iterable of possibly array-like arguments to check for
74        __array_function__ methods.
75
76    Returns
77    -------
78    Sequence of arguments with __array_function__ methods, in the order in
79    which they should be called.
80    """)
81
82
83ArgSpec = collections.namedtuple('ArgSpec', 'args varargs keywords defaults')
84
85
86def verify_matching_signatures(implementation, dispatcher):
87    """Verify that a dispatcher function has the right signature."""
88    implementation_spec = ArgSpec(*getargspec(implementation))
89    dispatcher_spec = ArgSpec(*getargspec(dispatcher))
90
91    if (implementation_spec.args != dispatcher_spec.args or
92            implementation_spec.varargs != dispatcher_spec.varargs or
93            implementation_spec.keywords != dispatcher_spec.keywords or
94            (bool(implementation_spec.defaults) !=
95             bool(dispatcher_spec.defaults)) or
96            (implementation_spec.defaults is not None and
97             len(implementation_spec.defaults) !=
98             len(dispatcher_spec.defaults))):
99        raise RuntimeError('implementation and dispatcher for %s have '
100                           'different function signatures' % implementation)
101
102    if implementation_spec.defaults is not None:
103        if dispatcher_spec.defaults != (None,) * len(dispatcher_spec.defaults):
104            raise RuntimeError('dispatcher functions can only use None for '
105                               'default argument values')
106
107
108def array_function_dispatch(dispatcher=None, module=None, verify=True,
109                            docs_from_dispatcher=False):
110    """Decorator for adding dispatch with the __array_function__ protocol.
111
112    See NEP-18 for example usage.
113
114    Parameters
115    ----------
116    dispatcher : callable or None
117        Function that when called like ``dispatcher(*args, **kwargs)`` with
118        arguments from the NumPy function call returns an iterable of
119        array-like arguments to check for ``__array_function__``.
120
121        If `None`, the first argument is used as the single `like=` argument
122        and not passed on.  A function implementing `like=` must call its
123        dispatcher with `like` as the first non-keyword argument.
124    module : str, optional
125        __module__ attribute to set on new function, e.g., ``module='numpy'``.
126        By default, module is copied from the decorated function.
127    verify : bool, optional
128        If True, verify the that the signature of the dispatcher and decorated
129        function signatures match exactly: all required and optional arguments
130        should appear in order with the same names, but the default values for
131        all optional arguments should be ``None``. Only disable verification
132        if the dispatcher's signature needs to deviate for some particular
133        reason, e.g., because the function has a signature like
134        ``func(*args, **kwargs)``.
135    docs_from_dispatcher : bool, optional
136        If True, copy docs from the dispatcher function onto the dispatched
137        function, rather than from the implementation. This is useful for
138        functions defined in C, which otherwise don't have docstrings.
139
140    Returns
141    -------
142    Function suitable for decorating the implementation of a NumPy function.
143
144    """
145    def decorator(implementation):
146        if verify:
147            if dispatcher is not None:
148                verify_matching_signatures(implementation, dispatcher)
149            else:
150                # Using __code__ directly similar to verify_matching_signature
151                co = implementation.__code__
152                last_arg = co.co_argcount + co.co_kwonlyargcount - 1
153                last_arg = co.co_varnames[last_arg]
154                if last_arg != "like" or co.co_kwonlyargcount == 0:
155                    raise RuntimeError(
156                        "__array_function__ expects `like=` to be the last "
157                        "argument and a keyword-only argument. "
158                        f"{implementation} does not seem to comply.")
159
160        if docs_from_dispatcher and dispatcher.__doc__ is not None:
161            doc = inspect.cleandoc(dispatcher.__doc__)
162            add_docstring(implementation, doc)
163
164        public_api = _ArrayFunctionDispatcher(dispatcher, implementation)
165        functools.update_wrapper(public_api, implementation)
166
167        if not verify and not getattr(implementation, "__text_signature__", None):
168            public_api.__signature__ = inspect.signature(dispatcher)
169
170        if module is not None:
171            public_api.__module__ = module
172
173        ARRAY_FUNCTIONS.add(public_api)
174
175        return public_api
176
177    return decorator
178
179
180def array_function_from_dispatcher(
181        implementation, module=None, verify=True, docs_from_dispatcher=True):
182    """Like array_function_dispatcher, but with function arguments flipped."""
183
184    def decorator(dispatcher):
185        return array_function_dispatch(
186            dispatcher, module, verify=verify,
187            docs_from_dispatcher=docs_from_dispatcher)(implementation)
188    return decorator
189