Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""Implementation of __array_function__ overrides from NEP-18."""
2import collections
3import functools
4import inspect
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
14ARRAY_FUNCTIONS = set()
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)
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)
30def finalize_array_function_like(public_api):
31 public_api.__doc__ = get_array_function_like_doc(public_api)
32 return public_api
35add_docstring(
36 _ArrayFunctionDispatcher,
37 """
38 Class to wrap functions with checks for __array_function__ overrides.
40 All arguments are required, and can only be passed by position.
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``.
57 Attributes
58 ----------
59 _implementation : function
60 The original implementation passed in.
61 """)
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__.
70 Parameters
71 ----------
72 relevant_args : iterable of array-like
73 Iterable of possibly array-like arguments to check for
74 __array_function__ methods.
76 Returns
77 -------
78 Sequence of arguments with __array_function__ methods, in the order in
79 which they should be called.
80 """)
83ArgSpec = collections.namedtuple('ArgSpec', 'args varargs keywords defaults')
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))
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)
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')
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.
112 See NEP-18 for example usage.
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__``.
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.
140 Returns
141 -------
142 Function suitable for decorating the implementation of a NumPy function.
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.")
160 if docs_from_dispatcher and dispatcher.__doc__ is not None:
161 doc = inspect.cleandoc(dispatcher.__doc__)
162 add_docstring(implementation, doc)
164 public_api = _ArrayFunctionDispatcher(dispatcher, implementation)
165 functools.update_wrapper(public_api, implementation)
167 if not verify and not getattr(implementation, "__text_signature__", None):
168 public_api.__signature__ = inspect.signature(dispatcher)
170 if module is not None:
171 public_api.__module__ = module
173 ARRAY_FUNCTIONS.add(public_api)
175 return public_api
177 return decorator
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."""
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