Coverage for /home/martinb/.local/share/virtualenvs/camcops/lib/python3.6/site-packages/numpy/core/overrides.py : 83%

Hot-keys on this page
r m x p 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 os
5import textwrap
7from numpy.core._multiarray_umath import (
8 add_docstring, implement_array_function, _get_implementing_args)
9from numpy.compat._inspect import getargspec
12ARRAY_FUNCTION_ENABLED = bool(
13 int(os.environ.get('NUMPY_EXPERIMENTAL_ARRAY_FUNCTION', 1)))
16add_docstring(
17 implement_array_function,
18 """
19 Implement a function with checks for __array_function__ overrides.
21 All arguments are required, and can only be passed by position.
23 Arguments
24 ---------
25 implementation : function
26 Function that implements the operation on NumPy array without
27 overrides when called like ``implementation(*args, **kwargs)``.
28 public_api : function
29 Function exposed by NumPy's public API originally called like
30 ``public_api(*args, **kwargs)`` on which arguments are now being
31 checked.
32 relevant_args : iterable
33 Iterable of arguments to check for __array_function__ methods.
34 args : tuple
35 Arbitrary positional arguments originally passed into ``public_api``.
36 kwargs : dict
37 Arbitrary keyword arguments originally passed into ``public_api``.
39 Returns
40 -------
41 Result from calling ``implementation()`` or an ``__array_function__``
42 method, as appropriate.
44 Raises
45 ------
46 TypeError : if no implementation is found.
47 """)
50# exposed for testing purposes; used internally by implement_array_function
51add_docstring(
52 _get_implementing_args,
53 """
54 Collect arguments on which to call __array_function__.
56 Parameters
57 ----------
58 relevant_args : iterable of array-like
59 Iterable of possibly array-like arguments to check for
60 __array_function__ methods.
62 Returns
63 -------
64 Sequence of arguments with __array_function__ methods, in the order in
65 which they should be called.
66 """)
69ArgSpec = collections.namedtuple('ArgSpec', 'args varargs keywords defaults')
72def verify_matching_signatures(implementation, dispatcher):
73 """Verify that a dispatcher function has the right signature."""
74 implementation_spec = ArgSpec(*getargspec(implementation))
75 dispatcher_spec = ArgSpec(*getargspec(dispatcher))
77 if (implementation_spec.args != dispatcher_spec.args or
78 implementation_spec.varargs != dispatcher_spec.varargs or
79 implementation_spec.keywords != dispatcher_spec.keywords or
80 (bool(implementation_spec.defaults) !=
81 bool(dispatcher_spec.defaults)) or
82 (implementation_spec.defaults is not None and
83 len(implementation_spec.defaults) !=
84 len(dispatcher_spec.defaults))):
85 raise RuntimeError('implementation and dispatcher for %s have '
86 'different function signatures' % implementation)
88 if implementation_spec.defaults is not None:
89 if dispatcher_spec.defaults != (None,) * len(dispatcher_spec.defaults):
90 raise RuntimeError('dispatcher functions can only use None for '
91 'default argument values')
94def set_module(module):
95 """Decorator for overriding __module__ on a function or class.
97 Example usage::
99 @set_module('numpy')
100 def example():
101 pass
103 assert example.__module__ == 'numpy'
104 """
105 def decorator(func):
106 if module is not None:
107 func.__module__ = module
108 return func
109 return decorator
113# Call textwrap.dedent here instead of in the function so as to avoid
114# calling dedent multiple times on the same text
115_wrapped_func_source = textwrap.dedent("""
116 @functools.wraps(implementation)
117 def {name}(*args, **kwargs):
118 relevant_args = dispatcher(*args, **kwargs)
119 return implement_array_function(
120 implementation, {name}, relevant_args, args, kwargs)
121 """)
124def array_function_dispatch(dispatcher, module=None, verify=True,
125 docs_from_dispatcher=False):
126 """Decorator for adding dispatch with the __array_function__ protocol.
128 See NEP-18 for example usage.
130 Parameters
131 ----------
132 dispatcher : callable
133 Function that when called like ``dispatcher(*args, **kwargs)`` with
134 arguments from the NumPy function call returns an iterable of
135 array-like arguments to check for ``__array_function__``.
136 module : str, optional
137 __module__ attribute to set on new function, e.g., ``module='numpy'``.
138 By default, module is copied from the decorated function.
139 verify : bool, optional
140 If True, verify the that the signature of the dispatcher and decorated
141 function signatures match exactly: all required and optional arguments
142 should appear in order with the same names, but the default values for
143 all optional arguments should be ``None``. Only disable verification
144 if the dispatcher's signature needs to deviate for some particular
145 reason, e.g., because the function has a signature like
146 ``func(*args, **kwargs)``.
147 docs_from_dispatcher : bool, optional
148 If True, copy docs from the dispatcher function onto the dispatched
149 function, rather than from the implementation. This is useful for
150 functions defined in C, which otherwise don't have docstrings.
152 Returns
153 -------
154 Function suitable for decorating the implementation of a NumPy function.
155 """
157 if not ARRAY_FUNCTION_ENABLED:
158 def decorator(implementation):
159 if docs_from_dispatcher:
160 add_docstring(implementation, dispatcher.__doc__)
161 if module is not None:
162 implementation.__module__ = module
163 return implementation
164 return decorator
166 def decorator(implementation):
167 if verify:
168 verify_matching_signatures(implementation, dispatcher)
170 if docs_from_dispatcher:
171 add_docstring(implementation, dispatcher.__doc__)
173 # Equivalently, we could define this function directly instead of using
174 # exec. This version has the advantage of giving the helper function a
175 # more interpettable name. Otherwise, the original function does not
176 # show up at all in many cases, e.g., if it's written in C or if the
177 # dispatcher gets an invalid keyword argument.
178 source = _wrapped_func_source.format(name=implementation.__name__)
180 source_object = compile(
181 source, filename='<__array_function__ internals>', mode='exec')
182 scope = {
183 'implementation': implementation,
184 'dispatcher': dispatcher,
185 'functools': functools,
186 'implement_array_function': implement_array_function,
187 }
188 exec(source_object, scope)
190 public_api = scope[implementation.__name__]
192 if module is not None:
193 public_api.__module__ = module
195 public_api._implementation = implementation
197 return public_api
199 return decorator
202def array_function_from_dispatcher(
203 implementation, module=None, verify=True, docs_from_dispatcher=True):
204 """Like array_function_dispatcher, but with function arguments flipped."""
206 def decorator(dispatcher):
207 return array_function_dispatch(
208 dispatcher, module, verify=verify,
209 docs_from_dispatcher=docs_from_dispatcher)(implementation)
210 return decorator