Enhancements to unraisable plugin (#12958)

A number of `unraisable` plugin enhancements:

* Set the unraisablehook as early as possible and unset it as late as possible, to collect the most possible unraisable exceptions (please let me know if there's an earlier or later hook I can use).
* Call the garbage collector just before unsetting the unraisablehook, to collect any straggling exceptions
* Collect multiple unraisable exceptions per test phase.
* Report the tracemalloc allocation traceback, if available.
* Avoid using a generator based hook to allow handling StopIteration in test failures
* Report the unraisable exception as the cause of the PytestUnraisableExceptionWarning exception if raised.
* Compute the repr of the unraisable.object in the unraisablehook so you get the latest information if available, and should help with resurrection of the object.

Co-authored-by: Ran Benita <[email protected]>
Co-authored-by: Bruno Oliveira <[email protected]>

Author: 3 people

changelog/12958.improvement.rst

@@ -0,0 +1,9 @@
+A number of :ref:`unraisable <unraisable>` enhancements:
+
+* Set the unraisable hook as early as possible and unset it as late as possible, to collect the most possible number of unraisable exceptions.
+* Call the garbage collector just before unsetting the unraisable hook, to collect any straggling exceptions.
+* Collect multiple unraisable exceptions per test phase.
+* Report the :mod:`tracemalloc` allocation traceback (if available).
+* Avoid using a generator based hook to allow handling :class:`StopIteration` in test failures.
+* Report the unraisable exception as the cause of the :class:`pytest.PytestUnraisableExceptionWarning` exception if raised.
+* Compute the ``repr`` of the unraisable object in the unraisable hook so you get the latest information if available, and should help with resurrection of the object.

src/_pytest/tracemalloc.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+
+def tracemalloc_message(source: object) -> str:
+    if source is None:
+        return ""
+
+    try:
+        import tracemalloc
+    except ImportError:
+        return ""
+
+    tb = tracemalloc.get_object_traceback(source)
+    if tb is not None:
+        formatted_tb = "\n".join(tb.format())
+        # Use a leading new line to better separate the (large) output
+        # from the traceback to the previous warning text.
+        return f"\nObject allocated at:\n{formatted_tb}"
+    # No need for a leading new line.
+    url = "https://docs.pytest.org/en/stable/how-to/capture-warnings.html#resource-warnings"
+    return (
+        "Enable tracemalloc to get traceback where the object was allocated.\n"
+        f"See {url} for more info."
+    )

src/_pytest/unraisableexception.py

@@ -1,100 +1,158 @@
 from __future__ import annotations
 
+import collections
 from collections.abc import Callable
-from collections.abc import Generator
+import functools
+import gc
 import sys
 import traceback
-from types import TracebackType
-from typing import Any
+from typing import NamedTuple
 from typing import TYPE_CHECKING
 import warnings
 
+from _pytest.config import Config
+from _pytest.nodes import Item
+from _pytest.stash import StashKey
+from _pytest.tracemalloc import tracemalloc_message
 import pytest
 
 
 if TYPE_CHECKING:
-    from typing_extensions import Self
-
-
-# Copied from cpython/Lib/test/support/__init__.py, with modifications.
-class catch_unraisable_exception:
-    """Context manager catching unraisable exception using sys.unraisablehook.
-
-    Storing the exception value (cm.unraisable.exc_value) creates a reference
-    cycle. The reference cycle is broken explicitly when the context manager
-    exits.
-
-    Storing the object (cm.unraisable.object) can resurrect it if it is set to
-    an object which is being finalized. Exiting the context manager clears the
-    stored object.
-
-    Usage:
-        with catch_unraisable_exception() as cm:
-            # code creating an "unraisable exception"
-            ...
-            # check the unraisable exception: use cm.unraisable
-            ...
-        # cm.unraisable attribute no longer exists at this point
-        # (to break a reference cycle)
-    """
-
-    def __init__(self) -> None:
-        self.unraisable: sys.UnraisableHookArgs | None = None
-        self._old_hook: Callable[[sys.UnraisableHookArgs], Any] | None = None
-
-    def _hook(self, unraisable: sys.UnraisableHookArgs) -> None:
-        # Storing unraisable.object can resurrect an object which is being
-        # finalized. Storing unraisable.exc_value creates a reference cycle.
-        self.unraisable = unraisable
-
-    def __enter__(self) -> Self:
-        self._old_hook = sys.unraisablehook
-        sys.unraisablehook = self._hook
-        return self
-
-    def __exit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_val: BaseException | None,
-        exc_tb: TracebackType | None,
-    ) -> None:
-        assert self._old_hook is not None
-        sys.unraisablehook = self._old_hook
-        self._old_hook = None
-        del self.unraisable
-
-
-def unraisable_exception_runtest_hook() -> Generator[None]:
-    with catch_unraisable_exception() as cm:
-        try:
-            yield
-        finally:
-            if cm.unraisable:
-                if cm.unraisable.err_msg is not None:
-                    err_msg = cm.unraisable.err_msg
-                else:
-                    err_msg = "Exception ignored in"
-                msg = f"{err_msg}: {cm.unraisable.object!r}\n\n"
-                msg += "".join(
-                    traceback.format_exception(
-                        cm.unraisable.exc_type,
-                        cm.unraisable.exc_value,
-                        cm.unraisable.exc_traceback,
-                    )
-                )
-                warnings.warn(pytest.PytestUnraisableExceptionWarning(msg))
+    pass
+
+if sys.version_info < (3, 11):
+    from exceptiongroup import ExceptionGroup
+
+
+def gc_collect_harder() -> None:
+    # A single collection doesn't necessarily collect everything.
+    # Constant determined experimentally by the Trio project.
+    for _ in range(5):
+        gc.collect()
 
 
-@pytest.hookimpl(wrapper=True, tryfirst=True)
-def pytest_runtest_setup() -> Generator[None]:
-    yield from unraisable_exception_runtest_hook()
+class UnraisableMeta(NamedTuple):
+    msg: str
+    cause_msg: str
+    exc_value: BaseException | None
 
 
-@pytest.hookimpl(wrapper=True, tryfirst=True)
-def pytest_runtest_call() -> Generator[None]:
-    yield from unraisable_exception_runtest_hook()
+unraisable_exceptions: StashKey[collections.deque[UnraisableMeta | BaseException]] = (
+    StashKey()
+)
 
 
-@pytest.hookimpl(wrapper=True, tryfirst=True)
-def pytest_runtest_teardown() -> Generator[None]:
-    yield from unraisable_exception_runtest_hook()
+def collect_unraisable(config: Config) -> None:
+    pop_unraisable = config.stash[unraisable_exceptions].pop
+    errors: list[pytest.PytestUnraisableExceptionWarning | RuntimeError] = []
+    meta = None
+    hook_error = None
+    try:
+        while True:
+            try:
+                meta = pop_unraisable()
+            except IndexError:
+                break
+
+            if isinstance(meta, BaseException):
+                hook_error = RuntimeError("Failed to process unraisable exception")
+                hook_error.__cause__ = meta
+                errors.append(hook_error)
+                continue
+
+            msg = meta.msg
+            try:
+                warnings.warn(pytest.PytestUnraisableExceptionWarning(msg))
+            except pytest.PytestUnraisableExceptionWarning as e:
+                # This except happens when the warning is treated as an error (e.g. `-Werror`).
+                if meta.exc_value is not None:
+                    # Exceptions have a better way to show the traceback, but
+                    # warnings do not, so hide the traceback from the msg and
+                    # set the cause so the traceback shows up in the right place.
+                    e.args = (meta.cause_msg,)
+                    e.__cause__ = meta.exc_value
+                errors.append(e)
+
+        if len(errors) == 1:
+            raise errors[0]
+        if errors:
+            raise ExceptionGroup("multiple unraisable exception warnings", errors)
+    finally:
+        del errors, meta, hook_error
+
+
+def cleanup(
+    *, config: Config, prev_hook: Callable[[sys.UnraisableHookArgs], object]
+) -> None:
+    try:
+        try:
+            gc_collect_harder()
+            collect_unraisable(config)
+        finally:
+            sys.unraisablehook = prev_hook
+    finally:
+        del config.stash[unraisable_exceptions]
+
+
+def unraisable_hook(
+    unraisable: sys.UnraisableHookArgs,
+    /,
+    *,
+    append: Callable[[UnraisableMeta | BaseException], object],
+) -> None:
+    try:
+        err_msg = (
+            "Exception ignored in" if unraisable.err_msg is None else unraisable.err_msg
+        )
+        summary = f"{err_msg}: {unraisable.object!r}"
+        traceback_message = "\n\n" + "".join(
+            traceback.format_exception(
+                unraisable.exc_type,
+                unraisable.exc_value,
+                unraisable.exc_traceback,
+            )
+        )
+        tracemalloc_tb = "\n" + tracemalloc_message(unraisable.object)
+        msg = summary + traceback_message + tracemalloc_tb
+        cause_msg = summary + tracemalloc_tb
+
+        append(
+            UnraisableMeta(
+                # we need to compute these strings here as they might change after
+                # the unraisablehook finishes and before the unraisable object is
+                # collected by a hook
+                msg=msg,
+                cause_msg=cause_msg,
+                exc_value=unraisable.exc_value,
+            )
+        )
+    except BaseException as e:
+        append(e)
+        # Raising this will cause the exception to be logged twice, once in our
+        # collect_unraisable and once by the unraisablehook calling machinery
+        # which is fine - this should never happen anyway and if it does
+        # it should probably be reported as a pytest bug.
+        raise
+
+
+def pytest_configure(config: Config) -> None:
+    prev_hook = sys.unraisablehook
+    deque: collections.deque[UnraisableMeta | BaseException] = collections.deque()
+    config.stash[unraisable_exceptions] = deque
+    config.add_cleanup(functools.partial(cleanup, config=config, prev_hook=prev_hook))
+    sys.unraisablehook = functools.partial(unraisable_hook, append=deque.append)
+
+
+@pytest.hookimpl(trylast=True)
+def pytest_runtest_setup(item: Item) -> None:
+    collect_unraisable(item.config)
+
+
+@pytest.hookimpl(trylast=True)
+def pytest_runtest_call(item: Item) -> None:
+    collect_unraisable(item.config)
+
+
+@pytest.hookimpl(trylast=True)
+def pytest_runtest_teardown(item: Item) -> None:
+    collect_unraisable(item.config)

src/_pytest/warnings.py

@@ -13,6 +13,7 @@
 from _pytest.main import Session
 from _pytest.nodes import Item
 from _pytest.terminal import TerminalReporter
+from _pytest.tracemalloc import tracemalloc_message
 import pytest
 
 
@@ -76,32 +77,13 @@ def catch_warnings_for_item(
 
 def warning_record_to_str(warning_message: warnings.WarningMessage) -> str:
     """Convert a warnings.WarningMessage to a string."""
-    warn_msg = warning_message.message
-    msg = warnings.formatwarning(
-        str(warn_msg),
+    return warnings.formatwarning(
+        str(warning_message.message),
         warning_message.category,
         warning_message.filename,
         warning_message.lineno,
         warning_message.line,
-    )
-    if warning_message.source is not None:
-        try:
-            import tracemalloc
-        except ImportError:
-            pass
-        else:
-            tb = tracemalloc.get_object_traceback(warning_message.source)
-            if tb is not None:
-                formatted_tb = "\n".join(tb.format())
-                # Use a leading new line to better separate the (large) output
-                # from the traceback to the previous warning text.
-                msg += f"\nObject allocated at:\n{formatted_tb}"
-            else:
-                # No need for a leading new line.
-                url = "https://docs.pytest.org/en/stable/how-to/capture-warnings.html#resource-warnings"
-                msg += "Enable tracemalloc to get traceback where the object was allocated.\n"
-                msg += f"See {url} for more info."
-    return msg
+    ) + tracemalloc_message(warning_message.source)
 
 
 @pytest.hookimpl(wrapper=True, tryfirst=True)