overrides.py 7.3 KB

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