threadexception.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152
  1. from __future__ import annotations
  2. import collections
  3. from collections.abc import Callable
  4. import functools
  5. import sys
  6. import threading
  7. import traceback
  8. from typing import NamedTuple
  9. from typing import TYPE_CHECKING
  10. import warnings
  11. from _pytest.config import Config
  12. from _pytest.nodes import Item
  13. from _pytest.stash import StashKey
  14. from _pytest.tracemalloc import tracemalloc_message
  15. import pytest
  16. if TYPE_CHECKING:
  17. pass
  18. if sys.version_info < (3, 11):
  19. from exceptiongroup import ExceptionGroup
  20. class ThreadExceptionMeta(NamedTuple):
  21. msg: str
  22. cause_msg: str
  23. exc_value: BaseException | None
  24. thread_exceptions: StashKey[collections.deque[ThreadExceptionMeta | BaseException]] = (
  25. StashKey()
  26. )
  27. def collect_thread_exception(config: Config) -> None:
  28. pop_thread_exception = config.stash[thread_exceptions].pop
  29. errors: list[pytest.PytestUnhandledThreadExceptionWarning | RuntimeError] = []
  30. meta = None
  31. hook_error = None
  32. try:
  33. while True:
  34. try:
  35. meta = pop_thread_exception()
  36. except IndexError:
  37. break
  38. if isinstance(meta, BaseException):
  39. hook_error = RuntimeError("Failed to process thread exception")
  40. hook_error.__cause__ = meta
  41. errors.append(hook_error)
  42. continue
  43. msg = meta.msg
  44. try:
  45. warnings.warn(pytest.PytestUnhandledThreadExceptionWarning(msg))
  46. except pytest.PytestUnhandledThreadExceptionWarning as e:
  47. # This except happens when the warning is treated as an error (e.g. `-Werror`).
  48. if meta.exc_value is not None:
  49. # Exceptions have a better way to show the traceback, but
  50. # warnings do not, so hide the traceback from the msg and
  51. # set the cause so the traceback shows up in the right place.
  52. e.args = (meta.cause_msg,)
  53. e.__cause__ = meta.exc_value
  54. errors.append(e)
  55. if len(errors) == 1:
  56. raise errors[0]
  57. if errors:
  58. raise ExceptionGroup("multiple thread exception warnings", errors)
  59. finally:
  60. del errors, meta, hook_error
  61. def cleanup(
  62. *, config: Config, prev_hook: Callable[[threading.ExceptHookArgs], object]
  63. ) -> None:
  64. try:
  65. try:
  66. # We don't join threads here, so exceptions raised from any
  67. # threads still running by the time _threading_atexits joins them
  68. # do not get captured (see #13027).
  69. collect_thread_exception(config)
  70. finally:
  71. threading.excepthook = prev_hook
  72. finally:
  73. del config.stash[thread_exceptions]
  74. def thread_exception_hook(
  75. args: threading.ExceptHookArgs,
  76. /,
  77. *,
  78. append: Callable[[ThreadExceptionMeta | BaseException], object],
  79. ) -> None:
  80. try:
  81. # we need to compute these strings here as they might change after
  82. # the excepthook finishes and before the metadata object is
  83. # collected by a pytest hook
  84. thread_name = "<unknown>" if args.thread is None else args.thread.name
  85. summary = f"Exception in thread {thread_name}"
  86. traceback_message = "\n\n" + "".join(
  87. traceback.format_exception(
  88. args.exc_type,
  89. args.exc_value,
  90. args.exc_traceback,
  91. )
  92. )
  93. tracemalloc_tb = "\n" + tracemalloc_message(args.thread)
  94. msg = summary + traceback_message + tracemalloc_tb
  95. cause_msg = summary + tracemalloc_tb
  96. append(
  97. ThreadExceptionMeta(
  98. # Compute these strings here as they might change later
  99. msg=msg,
  100. cause_msg=cause_msg,
  101. exc_value=args.exc_value,
  102. )
  103. )
  104. except BaseException as e:
  105. append(e)
  106. # Raising this will cause the exception to be logged twice, once in our
  107. # collect_thread_exception and once by sys.excepthook
  108. # which is fine - this should never happen anyway and if it does
  109. # it should probably be reported as a pytest bug.
  110. raise
  111. def pytest_configure(config: Config) -> None:
  112. prev_hook = threading.excepthook
  113. deque: collections.deque[ThreadExceptionMeta | BaseException] = collections.deque()
  114. config.stash[thread_exceptions] = deque
  115. config.add_cleanup(functools.partial(cleanup, config=config, prev_hook=prev_hook))
  116. threading.excepthook = functools.partial(thread_exception_hook, append=deque.append)
  117. @pytest.hookimpl(trylast=True)
  118. def pytest_runtest_setup(item: Item) -> None:
  119. collect_thread_exception(item.config)
  120. @pytest.hookimpl(trylast=True)
  121. def pytest_runtest_call(item: Item) -> None:
  122. collect_thread_exception(item.config)
  123. @pytest.hookimpl(trylast=True)
  124. def pytest_runtest_teardown(item: Item) -> None:
  125. collect_thread_exception(item.config)