_callers.py 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155
  1. """
  2. Call loop machinery
  3. """
  4. from __future__ import annotations
  5. from typing import cast
  6. from typing import Generator
  7. from typing import Mapping
  8. from typing import Sequence
  9. from typing import Tuple
  10. from typing import TYPE_CHECKING
  11. from typing import Union
  12. from ._result import _raise_wrapfail
  13. from ._result import _Result
  14. from ._result import HookCallError
  15. if TYPE_CHECKING:
  16. from ._hooks import HookImpl
  17. # Need to distinguish between old- and new-style hook wrappers.
  18. # Wrapping one a singleton tuple is the fastest type-safe way I found to do it.
  19. Teardown = Union[
  20. Tuple[Generator[None, _Result[object], None]],
  21. Generator[None, object, object],
  22. ]
  23. def _multicall(
  24. hook_name: str,
  25. hook_impls: Sequence[HookImpl],
  26. caller_kwargs: Mapping[str, object],
  27. firstresult: bool,
  28. ) -> object | list[object]:
  29. """Execute a call into multiple python functions/methods and return the
  30. result(s).
  31. ``caller_kwargs`` comes from _HookCaller.__call__().
  32. """
  33. __tracebackhide__ = True
  34. results: list[object] = []
  35. exception = None
  36. only_new_style_wrappers = True
  37. try: # run impl and wrapper setup functions in a loop
  38. teardowns: list[Teardown] = []
  39. try:
  40. for hook_impl in reversed(hook_impls):
  41. try:
  42. args = [caller_kwargs[argname] for argname in hook_impl.argnames]
  43. except KeyError:
  44. for argname in hook_impl.argnames:
  45. if argname not in caller_kwargs:
  46. raise HookCallError(
  47. f"hook call must provide argument {argname!r}"
  48. )
  49. if hook_impl.hookwrapper:
  50. only_new_style_wrappers = False
  51. try:
  52. # If this cast is not valid, a type error is raised below,
  53. # which is the desired response.
  54. res = hook_impl.function(*args)
  55. wrapper_gen = cast(Generator[None, _Result[object], None], res)
  56. next(wrapper_gen) # first yield
  57. teardowns.append((wrapper_gen,))
  58. except StopIteration:
  59. _raise_wrapfail(wrapper_gen, "did not yield")
  60. elif hook_impl.wrapper:
  61. try:
  62. # If this cast is not valid, a type error is raised below,
  63. # which is the desired response.
  64. res = hook_impl.function(*args)
  65. function_gen = cast(Generator[None, object, object], res)
  66. next(function_gen) # first yield
  67. teardowns.append(function_gen)
  68. except StopIteration:
  69. _raise_wrapfail(function_gen, "did not yield")
  70. else:
  71. res = hook_impl.function(*args)
  72. if res is not None:
  73. results.append(res)
  74. if firstresult: # halt further impl calls
  75. break
  76. except BaseException as exc:
  77. exception = exc
  78. finally:
  79. # Fast path - only new-style wrappers, no _Result.
  80. if only_new_style_wrappers:
  81. if firstresult: # first result hooks return a single value
  82. result = results[0] if results else None
  83. else:
  84. result = results
  85. # run all wrapper post-yield blocks
  86. for teardown in reversed(teardowns):
  87. try:
  88. if exception is not None:
  89. teardown.throw(exception) # type: ignore[union-attr]
  90. else:
  91. teardown.send(result) # type: ignore[union-attr]
  92. # Following is unreachable for a well behaved hook wrapper.
  93. # Try to force finalizers otherwise postponed till GC action.
  94. # Note: close() may raise if generator handles GeneratorExit.
  95. teardown.close() # type: ignore[union-attr]
  96. except StopIteration as si:
  97. result = si.value
  98. exception = None
  99. continue
  100. except BaseException as e:
  101. exception = e
  102. continue
  103. _raise_wrapfail(teardown, "has second yield") # type: ignore[arg-type]
  104. if exception is not None:
  105. raise exception.with_traceback(exception.__traceback__)
  106. else:
  107. return result
  108. # Slow path - need to support old-style wrappers.
  109. else:
  110. if firstresult: # first result hooks return a single value
  111. outcome: _Result[object | list[object]] = _Result(
  112. results[0] if results else None, exception
  113. )
  114. else:
  115. outcome = _Result(results, exception)
  116. # run all wrapper post-yield blocks
  117. for teardown in reversed(teardowns):
  118. if isinstance(teardown, tuple):
  119. try:
  120. teardown[0].send(outcome)
  121. _raise_wrapfail(teardown[0], "has second yield")
  122. except StopIteration:
  123. pass
  124. else:
  125. try:
  126. if outcome._exception is not None:
  127. teardown.throw(outcome._exception)
  128. else:
  129. teardown.send(outcome._result)
  130. # Following is unreachable for a well behaved hook wrapper.
  131. # Try to force finalizers otherwise postponed till GC action.
  132. # Note: close() may raise if generator handles GeneratorExit.
  133. teardown.close()
  134. except StopIteration as si:
  135. outcome.force_result(si.value)
  136. continue
  137. except BaseException as e:
  138. outcome.force_exception(e)
  139. continue
  140. _raise_wrapfail(teardown, "has second yield")
  141. return outcome.get_result()