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