no-copy solution that completely hides the exceptiongroup in most cases

Author: jakkdl

tests/test_connection.py

@@ -452,7 +452,6 @@ async def test_open_websocket_internal_ki(nursery, monkeypatch, autojump_clock):
     Make sure that KI is delivered, and the user exception is in the __cause__ exceptiongroup
     """
     async def ki_raising_ping_handler(*args, **kwargs) -> None:
-        print("raising ki")
         raise KeyboardInterrupt
     monkeypatch.setattr(WebSocketConnection, "_handle_ping_event", ki_raising_ping_handler)
     async def handler(request):
@@ -474,27 +473,32 @@ async def handler(request):
 async def test_open_websocket_internal_exc(nursery, monkeypatch, autojump_clock):
     """_reader_task._handle_ping_event triggers ValueError.
     user code also raises exception.
-    internal exception is in __cause__ exceptiongroup and user exc is delivered
+    internal exception is in __context__ exceptiongroup and user exc is delivered
     """
-    my_value_error = ValueError()
+    internal_error = ValueError()
+    internal_error.__context__ = TypeError()
+    user_error = NameError()
+    user_error_context = KeyError()
     async def raising_ping_event(*args, **kwargs) -> None:
-        raise my_value_error
+        raise internal_error
 
     monkeypatch.setattr(WebSocketConnection, "_handle_ping_event", raising_ping_event)
     async def handler(request):
         server_ws = await request.accept()
         await server_ws.ping(b"a")
 
     server = await nursery.start(serve_websocket, handler, HOST, 0, None)
-    with pytest.raises(trio.TooSlowError) as exc_info:
+    with pytest.raises(type(user_error)) as exc_info:
         async with open_websocket(HOST, server.port, RESOURCE, use_ssl=False):
-            with trio.fail_after(1) as cs:
-                cs.shield = True
-                await trio.sleep(2)
+            await trio.lowlevel.checkpoint()
+            user_error.__context__ = user_error_context
+            raise user_error
 
-    e_cause = exc_info.value.__cause__
-    assert isinstance(e_cause, _TRIO_EXC_GROUP_TYPE)
-    assert my_value_error in e_cause.exceptions
+    assert exc_info.value is user_error
+    e_context = exc_info.value.__context__
+    assert isinstance(e_context, BaseExceptionGroup)
+    assert internal_error in e_context.exceptions
+    assert user_error_context in e_context.exceptions
 
 @fail_after(5)
 async def test_open_websocket_cancellations(nursery, monkeypatch, autojump_clock):
@@ -528,14 +532,9 @@ async def handler(request):
                     user_cancelled_context = e.__context__
                     raise
 
-    # a copy of user_cancelled is reraised
-    assert exc_info.value is not user_cancelled
-    # with the same cause
+    assert exc_info.value is user_cancelled
     assert exc_info.value.__cause__ is user_cancelled_cause
-    # the context is the exception group, which contains the original user_cancelled
-    assert exc_info.value.__context__.exceptions[1] is user_cancelled
-    assert exc_info.value.__context__.exceptions[1].__cause__ is user_cancelled_cause
-    assert exc_info.value.__context__.exceptions[1].__context__ is user_cancelled_context
+    assert exc_info.value.__context__ is user_cancelled_context
 
 def _trio_default_non_strict_exception_groups() -> bool:
     assert re.match(r'^0\.\d\d\.', trio.__version__), "unexpected trio versioning scheme"
@@ -586,19 +585,9 @@ async def handler(request):
             except TypeError:
                 raise e_primary from e_cause
     e = exc_info.value
-    if _trio_default_non_strict_exception_groups():
-        assert e is e_primary
-        assert e.__cause__ is e_cause
-        assert e.__context__ is e_context
-    else:
-        # a copy is reraised to avoid losing e_context
-        assert e is not e_primary
-        assert e.__cause__ is e_cause
-
-        # the nursery-internal group is injected as context
-        assert isinstance(e.__context__, _TRIO_EXC_GROUP_TYPE)
-        assert e.__context__.exceptions[0] is e_primary
-        assert e.__context__.exceptions[0].__context__ is e_context
+    assert e is e_primary
+    assert e.__cause__ is e_cause
+    assert e.__context__ is e_context
 
 @fail_after(1)
 async def test_reject_handshake(nursery):

trio_websocket/_impl.py

@@ -12,7 +12,7 @@
 import ssl
 import struct
 import urllib.parse
-from typing import Iterable, List, Optional, Union
+from typing import Iterable, List, NoReturn, Optional, Union
 
 import outcome
 import trio
@@ -192,10 +192,29 @@ async def _close_connection(connection: WebSocketConnection) -> None:
         except trio.TooSlowError:
             raise DisconnectionTimeout from None
 
+    def _raise(exc: BaseException) -> NoReturn:
+        __tracebackhide__ = True
+        context = exc.__context__
+        try:
+            raise exc
+        finally:
+            exc.__context__ = context
+            del exc, context
+
     connection: WebSocketConnection|None=None
     close_result: outcome.Maybe[None] | None = None
     user_error = None
 
+    # Unwrapping exception groups has a lot of pitfalls, one of them stemming from
+    # the exception we raise also being inside the group that's set as the context.
+    # This leads to loss of info unless properly handled.
+    # See https://github.com/python-trio/flake8-async/issues/298
+    # We therefore save the exception before raising it, and save our intended context,
+    # so they can be modified in the `finally`.
+    exc_to_raise = None
+    exc_context = None
+    # by avoiding use of `raise .. from ..` we leave the original __cause__
+
     try:
         async with trio.open_nursery() as new_nursery:
             result = await outcome.acapture(_open_connection, new_nursery)
@@ -216,7 +235,7 @@ async def _close_connection(connection: WebSocketConnection) -> None:
     except _TRIO_EXC_GROUP_TYPE as e:
         # user_error, or exception bubbling up from _reader_task
         if len(e.exceptions) == 1:
-            raise copy_exc(e.exceptions[0]) from e.exceptions[0].__cause__
+            _raise(e.exceptions[0])
 
         # contains at most 1 non-cancelled exceptions
         exception_to_raise: BaseException|None = None
@@ -229,25 +248,40 @@ async def _close_connection(connection: WebSocketConnection) -> None:
         else:
             if exception_to_raise is None:
                 # all exceptions are cancelled
-                # prefer raising the one from the user, for traceback reasons
+                # we reraise the user exception and throw out internal
                 if user_error is not None:
-                    # no reason to raise from e, just to include a bunch of extra
-                    # cancelleds.
-                    raise copy_exc(user_error) from user_error.__cause__
+                    _raise(user_error)
                 # multiple internal Cancelled is not possible afaik
-                raise copy_exc(e.exceptions[0]) from e  # pragma: no cover
-            raise copy_exc(exception_to_raise) from exception_to_raise.__cause__
+                # but if so we just raise one of them
+                _raise(e.exceptions[0])
+            # raise the non-cancelled exception
+            _raise(exception_to_raise)
 
-        # if we have any KeyboardInterrupt in the group, make sure to raise it.
+        # if we have any KeyboardInterrupt in the group, raise a new KeyboardInterrupt
+        # with the group as cause & context
         for sub_exc in e.exceptions:
             if isinstance(sub_exc, KeyboardInterrupt):
-                raise copy_exc(sub_exc) from e
+                raise KeyboardInterrupt from e
 
         # Both user code and internal code raised non-cancelled exceptions.
-        # We "hide" the internal exception(s) in the __cause__ and surface
-        # the user_error.
+        # We set the context to be an exception group containing internal exceptions
+        # and, if not None, `user_error.__context__`
         if user_error is not None:
-            raise copy_exc(user_error) from e
+            exceptions = [subexc for subexc in e.exceptions if subexc is not user_error]
+            eg_substr = ''
+            # there's technically loss of info here, with __suppress_context__=True you
+            # still have original __context__ available, just not printed. But we delete
+            # it completely because we can't partially suppress the group
+            if user_error.__context__ is not None and not user_error.__suppress_context__:
+                exceptions.append(user_error.__context__)
+                eg_substr = ' and the context for the user exception'
+            eg_str = (
+                "Both internal and user exceptions encountered. This group contains "
+                "the internal exception(s)" + eg_substr + "."
+            )
+            user_error.__context__ = BaseExceptionGroup(eg_str, exceptions)
+            user_error.__suppress_context__ = False
+            _raise(user_error)
 
         raise TrioWebsocketInternalError(
             "The trio-websocket API is not expected to raise multiple exceptions. "