diff --git a/src/trio/_tests/test_util.py b/src/trio/_tests/test_util.py index 5036d76e52..ba11f5a311 100644 --- a/src/trio/_tests/test_util.py +++ b/src/trio/_tests/test_util.py @@ -24,9 +24,13 @@ fixup_module_metadata, generic_function, is_main_thread, + raise_single_exception_from_group, ) from ..testing import wait_all_tasks_blocked +if sys.version_info < (3, 11): + from exceptiongroup import BaseExceptionGroup, ExceptionGroup + if TYPE_CHECKING: from collections.abc import AsyncGenerator @@ -267,3 +271,68 @@ def test_fixup_module_metadata() -> None: mod.some_func() mod._private() mod.SomeClass().method() + + +async def test_raise_single_exception_from_group() -> None: + excinfo: pytest.ExceptionInfo[BaseException] + + exc = ValueError("foo") + cause = SyntaxError("cause") + context = TypeError("context") + exc.__cause__ = cause + exc.__context__ = context + cancelled = trio.Cancelled._create() + + with pytest.raises(ValueError, match="foo") as excinfo: + raise_single_exception_from_group(ExceptionGroup("", [exc])) + assert excinfo.value.__cause__ == cause + assert excinfo.value.__context__ == context + + with pytest.raises(ValueError, match="foo") as excinfo: + raise_single_exception_from_group( + ExceptionGroup("", [ExceptionGroup("", [exc])]) + ) + assert excinfo.value.__cause__ == cause + assert excinfo.value.__context__ == context + + with pytest.raises(ValueError, match="foo") as excinfo: + raise_single_exception_from_group( + BaseExceptionGroup( + "", [cancelled, BaseExceptionGroup("", [cancelled, exc])] + ) + ) + assert excinfo.value.__cause__ == cause + assert excinfo.value.__context__ == context + + # multiple non-cancelled + eg = ExceptionGroup("", [ValueError("foo"), ValueError("bar")]) + with pytest.raises( + AssertionError, + match=r"^Attempted to unwrap exceptiongroup with multiple non-cancelled exceptions. This is often caused by a bug in the caller.$", + ) as excinfo: + raise_single_exception_from_group(eg) + assert excinfo.value.__cause__ is eg + assert excinfo.value.__context__ is None + + # keyboardinterrupt overrides everything + eg_ki = BaseExceptionGroup( + "", + [ + ValueError("foo"), + ValueError("bar"), + KeyboardInterrupt("this exc doesn't get reraised"), + ], + ) + with pytest.raises(KeyboardInterrupt, match=r"^$") as excinfo: + raise_single_exception_from_group(eg_ki) + assert excinfo.value.__cause__ is eg_ki + assert excinfo.value.__context__ is None + + # if we only got cancelled, first one is reraised + with pytest.raises(trio.Cancelled, match=r"^Cancelled$") as excinfo: + raise_single_exception_from_group( + BaseExceptionGroup("", [cancelled, trio.Cancelled._create()]) + ) + assert excinfo.value is cancelled + assert excinfo.value.__cause__ is None + assert excinfo.value.__context__ is None diff --git a/src/trio/_util.py b/src/trio/_util.py index 9b8b1d7a46..54fc5ff733 100644 --- a/src/trio/_util.py +++ b/src/trio/_util.py @@ -4,6 +4,7 @@ import collections.abc import inspect import signal +import sys from abc import ABCMeta from collections.abc import Awaitable, Callable, Sequence from functools import update_wrapper @@ -20,6 +21,9 @@ import trio +if sys.version_info < (3, 11): + from exceptiongroup import BaseExceptionGroup + # Explicit "Any" is not allowed CallT = TypeVar("CallT", bound=Callable[..., Any]) # type: ignore[explicit-any] T = TypeVar("T") @@ -353,3 +357,61 @@ def wraps( # type: ignore[explicit-any] else: from functools import wraps # noqa: F401 # this is re-exported + + +def _raise(exc: BaseException) -> NoReturn: + """This helper allows re-raising an exception without __context__ being set.""" + # cause does not need special handling, we simply avoid using `raise .. from ..` + __tracebackhide__ = True + context = exc.__context__ + try: + raise exc + finally: + exc.__context__ = context + del exc, context + + +def raise_single_exception_from_group( + eg: BaseExceptionGroup[BaseException], +) -> NoReturn: + """This function takes an exception group that is assumed to have at most + one non-cancelled exception, which it reraises as a standalone exception. + + If a :exc:`KeyboardInterrupt` is encountered, a new KeyboardInterrupt is immediately + raised with the entire group as cause. + + If the group only contains :exc:`Cancelled` it reraises the first one encountered. + + It will retain context and cause of the contained exception, and entirely discard + the cause/context of the group(s). + + If multiple non-cancelled exceptions are encountered, it raises + :exc:`AssertionError`. + """ + cancelled_exceptions = [] + noncancelled_exceptions = [] + + # subgroup/split retains excgroup structure, so we need to manually traverse + def _parse_excg(e: BaseException) -> None: + if isinstance(e, (KeyboardInterrupt, SystemExit)): + # immediately bail out + raise KeyboardInterrupt from eg + + if isinstance(e, trio.Cancelled): + cancelled_exceptions.append(e) + elif isinstance(e, BaseExceptionGroup): + for sub_e in e.exceptions: + _parse_excg(sub_e) + else: + noncancelled_exceptions.append(e) + + _parse_excg(eg) + + if len(noncancelled_exceptions) > 1: + raise AssertionError( + "Attempted to unwrap exceptiongroup with multiple non-cancelled exceptions. This is often caused by a bug in the caller." + ) from eg + if len(noncancelled_exceptions) == 1: + _raise(noncancelled_exceptions[0]) + assert cancelled_exceptions, "internal error" + _raise(cancelled_exceptions[0]) diff --git a/tox.ini b/tox.ini index 46b9fae2bc..6eecf7b78a 100644 --- a/tox.ini +++ b/tox.ini @@ -17,7 +17,7 @@ labels = # * move to pyproject.toml? # * this means conditional deps need to be replaced -# protip: install uv-tox for faster venv generation +# protip: install tox-uv for faster venv generation [testenv] description = "Base environment for running tests depending on python version."