unraisablehook enhancements

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,131 @@
 from __future__ import annotations
 
+import collections
+import functools
+import gc
 import sys
 import traceback
-from types import TracebackType
-from typing import Any
 from typing import Callable
-from typing import Generator
+from typing import NamedTuple
 from typing import TYPE_CHECKING
 import warnings
 
+from _pytest.config import Config
+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
+
 
+class UnraisableMeta(NamedTuple):
+    msg: str
+    cause_msg: str
+    exc_value: BaseException | None
 
-@pytest.hookimpl(wrapper=True, tryfirst=True)
-def pytest_runtest_setup() -> Generator[None]:
-    yield from unraisable_exception_runtest_hook()
 
+_unraisable_exceptions: collections.deque[UnraisableMeta | BaseException] = (
+    collections.deque()
+)
 
-@pytest.hookimpl(wrapper=True, tryfirst=True)
-def pytest_runtest_call() -> Generator[None]:
-    yield from unraisable_exception_runtest_hook()
 
+def collect_unraisable() -> None:
+    errors: list[pytest.PytestUnraisableExceptionWarning | RuntimeError] = []
+    meta = None
+    hook_error = None
+    try:
+        while True:
+            try:
+                meta = _unraisable_exceptions.pop()
+            except IndexError:
+                break
 
-@pytest.hookimpl(wrapper=True, tryfirst=True)
-def pytest_runtest_teardown() -> Generator[None]:
-    yield from unraisable_exception_runtest_hook()
+            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:
+                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(prev_hook: Callable[[sys.UnraisableHookArgs], object]) -> None:
+    try:
+        for i in range(5):
+            gc.collect()
+        collect_unraisable()
+    finally:
+        sys.unraisablehook = prev_hook
+
+
+def unraisable_hook(unraisable: sys.UnraisableHookArgs) -> 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 = tracemalloc_message(unraisable.object)
+        msg = summary + traceback_message + tracemalloc_tb
+        cause_msg = summary + tracemalloc_tb
+
+        _unraisable_exceptions.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:
+        _unraisable_exceptions.append(e)
+
+
+def pytest_configure(config: Config) -> None:
+    prev_hook = sys.unraisablehook
+    config.add_cleanup(functools.partial(_cleanup, prev_hook))
+    sys.unraisablehook = unraisable_hook
+
+
+@pytest.hookimpl(trylast=True)
+def pytest_runtest_setup() -> None:
+    collect_unraisable()
+
+
+@pytest.hookimpl(trylast=True)
+def pytest_runtest_call() -> None:
+    collect_unraisable()
+
+
+@pytest.hookimpl(trylast=True)
+def pytest_runtest_teardown() -> None:
+    collect_unraisable()

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)