overrides.py 8.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231
  1. """Implementation of __array_function__ overrides from NEP-18."""
  2. import collections
  3. import functools
  4. import os
  5. import textwrap
  6. from numpy.core._multiarray_umath import (
  7. add_docstring, implement_array_function, _get_implementing_args)
  8. from numpy.compat._inspect import getargspec
  9. ARRAY_FUNCTION_ENABLED = bool(
  10. int(os.environ.get('NUMPY_EXPERIMENTAL_ARRAY_FUNCTION', 1)))
  11. array_function_like_doc = (
  12. """like : array_like
  13. Reference object to allow the creation of arrays which are not
  14. NumPy arrays. If an array-like passed in as ``like`` supports
  15. the ``__array_function__`` protocol, the result will be defined
  16. by it. In this case, it ensures the creation of an array object
  17. compatible with that passed in via this argument.
  18. .. note::
  19. The ``like`` keyword is an experimental feature pending on
  20. acceptance of :ref:`NEP 35 <NEP35>`."""
  21. )
  22. def set_array_function_like_doc(public_api):
  23. if public_api.__doc__ is not None:
  24. public_api.__doc__ = public_api.__doc__.replace(
  25. "${ARRAY_FUNCTION_LIKE}",
  26. array_function_like_doc,
  27. )
  28. return public_api
  29. add_docstring(
  30. implement_array_function,
  31. """
  32. Implement a function with checks for __array_function__ overrides.
  33. All arguments are required, and can only be passed by position.
  34. Parameters
  35. ----------
  36. implementation : function
  37. Function that implements the operation on NumPy array without
  38. overrides when called like ``implementation(*args, **kwargs)``.
  39. public_api : function
  40. Function exposed by NumPy's public API originally called like
  41. ``public_api(*args, **kwargs)`` on which arguments are now being
  42. checked.
  43. relevant_args : iterable
  44. Iterable of arguments to check for __array_function__ methods.
  45. args : tuple
  46. Arbitrary positional arguments originally passed into ``public_api``.
  47. kwargs : dict
  48. Arbitrary keyword arguments originally passed into ``public_api``.
  49. Returns
  50. -------
  51. Result from calling ``implementation()`` or an ``__array_function__``
  52. method, as appropriate.
  53. Raises
  54. ------
  55. TypeError : if no implementation is found.
  56. """)
  57. # exposed for testing purposes; used internally by implement_array_function
  58. add_docstring(
  59. _get_implementing_args,
  60. """
  61. Collect arguments on which to call __array_function__.
  62. Parameters
  63. ----------
  64. relevant_args : iterable of array-like
  65. Iterable of possibly array-like arguments to check for
  66. __array_function__ methods.
  67. Returns
  68. -------
  69. Sequence of arguments with __array_function__ methods, in the order in
  70. which they should be called.
  71. """)
  72. ArgSpec = collections.namedtuple('ArgSpec', 'args varargs keywords defaults')
  73. def verify_matching_signatures(implementation, dispatcher):
  74. """Verify that a dispatcher function has the right signature."""
  75. implementation_spec = ArgSpec(*getargspec(implementation))
  76. 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)
  87. if implementation_spec.defaults is not None:
  88. if dispatcher_spec.defaults != (None,) * len(dispatcher_spec.defaults):
  89. raise RuntimeError('dispatcher functions can only use None for '
  90. 'default argument values')
  91. def set_module(module):
  92. """Decorator for overriding __module__ on a function or class.
  93. Example usage::
  94. @set_module('numpy')
  95. def example():
  96. pass
  97. assert example.__module__ == 'numpy'
  98. """
  99. def decorator(func):
  100. if module is not None:
  101. func.__module__ = module
  102. return func
  103. return decorator
  104. # Call textwrap.dedent here instead of in the function so as to avoid
  105. # calling dedent multiple times on the same text
  106. _wrapped_func_source = textwrap.dedent("""
  107. @functools.wraps(implementation)
  108. def {name}(*args, **kwargs):
  109. relevant_args = dispatcher(*args, **kwargs)
  110. return implement_array_function(
  111. implementation, {name}, relevant_args, args, kwargs)
  112. """)
  113. def array_function_dispatch(dispatcher, module=None, verify=True,
  114. docs_from_dispatcher=False):
  115. """Decorator for adding dispatch with the __array_function__ protocol.
  116. See NEP-18 for example usage.
  117. Parameters
  118. ----------
  119. dispatcher : callable
  120. Function that when called like ``dispatcher(*args, **kwargs)`` with
  121. arguments from the NumPy function call returns an iterable of
  122. array-like arguments to check for ``__array_function__``.
  123. module : str, optional
  124. __module__ attribute to set on new function, e.g., ``module='numpy'``.
  125. By default, module is copied from the decorated function.
  126. verify : bool, optional
  127. If True, verify the that the signature of the dispatcher and decorated
  128. function signatures match exactly: all required and optional arguments
  129. should appear in order with the same names, but the default values for
  130. all optional arguments should be ``None``. Only disable verification
  131. if the dispatcher's signature needs to deviate for some particular
  132. reason, e.g., because the function has a signature like
  133. ``func(*args, **kwargs)``.
  134. docs_from_dispatcher : bool, optional
  135. If True, copy docs from the dispatcher function onto the dispatched
  136. function, rather than from the implementation. This is useful for
  137. functions defined in C, which otherwise don't have docstrings.
  138. Returns
  139. -------
  140. Function suitable for decorating the implementation of a NumPy function.
  141. """
  142. if not ARRAY_FUNCTION_ENABLED:
  143. def decorator(implementation):
  144. if docs_from_dispatcher:
  145. add_docstring(implementation, dispatcher.__doc__)
  146. if module is not None:
  147. implementation.__module__ = module
  148. return implementation
  149. return decorator
  150. def decorator(implementation):
  151. if verify:
  152. verify_matching_signatures(implementation, dispatcher)
  153. if docs_from_dispatcher:
  154. add_docstring(implementation, dispatcher.__doc__)
  155. # Equivalently, we could define this function directly instead of using
  156. # exec. This version has the advantage of giving the helper function a
  157. # more interpettable name. Otherwise, the original function does not
  158. # show up at all in many cases, e.g., if it's written in C or if the
  159. # dispatcher gets an invalid keyword argument.
  160. source = _wrapped_func_source.format(name=implementation.__name__)
  161. source_object = compile(
  162. source, filename='<__array_function__ internals>', mode='exec')
  163. scope = {
  164. 'implementation': implementation,
  165. 'dispatcher': dispatcher,
  166. 'functools': functools,
  167. 'implement_array_function': implement_array_function,
  168. }
  169. exec(source_object, scope)
  170. public_api = scope[implementation.__name__]
  171. if module is not None:
  172. public_api.__module__ = module
  173. public_api._implementation = implementation
  174. return public_api
  175. return decorator
  176. def array_function_from_dispatcher(
  177. implementation, module=None, verify=True, docs_from_dispatcher=True):
  178. """Like array_function_dispatcher, but with function arguments flipped."""
  179. def decorator(dispatcher):
  180. return array_function_dispatch(
  181. dispatcher, module, verify=verify,
  182. docs_from_dispatcher=docs_from_dispatcher)(implementation)
  183. return decorator