add helper for unwrapping exception groups with single exception (#3240)

* add raise_single_exception_from_group

Author: jakkdl

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

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])

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."