fix loss of context/cause on exceptions raised inside open_websocket

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)  # pylint: disable=possibly-used-before-assignment
+    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):
@@ -513,6 +517,8 @@ async def handler(request):
         server_ws = await request.accept()
         await server_ws.ping(b"a")
     user_cancelled = None
+    user_cancelled_cause = None
+    user_cancelled_context = None
 
     server = await nursery.start(serve_websocket, handler, HOST, 0, None)
     with trio.move_on_after(2):
@@ -522,8 +528,13 @@ async def handler(request):
                     await trio.sleep_forever()
                 except trio.Cancelled as e:
                     user_cancelled = e
+                    user_cancelled_cause = e.__cause__
+                    user_cancelled_context = e.__context__
                     raise
+
     assert exc_info.value is user_cancelled
+    assert exc_info.value.__cause__ is user_cancelled_cause
+    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"
@@ -560,6 +571,24 @@ async def handler(request):
                     RaisesGroup(ValueError)))).matches(exc.value)
 
 
+async def test_user_exception_cause(nursery) -> None:
+    async def handler(request):
+        await request.accept()
+    server = await nursery.start(serve_websocket, handler, HOST, 0, None)
+    e_context = TypeError("foo")
+    e_primary = ValueError("bar")
+    e_cause = RuntimeError("zee")
+    with pytest.raises(ValueError) as exc_info:
+        async with open_websocket(HOST, server.port, RESOURCE, use_ssl=False):
+            try:
+                raise e_context
+            except TypeError:
+                raise e_primary from e_cause
+    e = exc_info.value
+    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):
     async def handler(request):

trio_websocket/_impl.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 import sys
 from collections import OrderedDict
 from contextlib import asynccontextmanager
@@ -11,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
@@ -91,6 +92,16 @@ def __exit__(self, ty, value, tb):
             filtered_exception = _ignore_cancel(value)
         return filtered_exception is None
 
+def copy_exc(e: BaseException) -> BaseException:
+    """Copy an exception.
+
+    `copy.copy` fails on `trio.Cancelled`, and on exceptions with a custom `__init__`
+    that calls `super().__init__()`. It may be the case that this also fails on something.
+    """
+    cls = type(e)
+    result = cls.__new__(cls)
+    result.__dict__ = copy.copy(e.__dict__)
+    return result
 
 @asynccontextmanager
 async def open_websocket(
@@ -181,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)
@@ -205,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 e.exceptions[0]
+            _raise(e.exceptions[0])
 
         # contains at most 1 non-cancelled exceptions
         exception_to_raise: BaseException|None = None
@@ -218,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 user_error  # pylint: disable=raise-missing-from
+                    _raise(user_error)
                 # multiple internal Cancelled is not possible afaik
-                raise e.exceptions[0]  # pragma: no cover  # pylint: disable=raise-missing-from
-            raise exception_to_raise
+                # 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 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 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. "