Correctly preserve exception __context__ in MultiError.catch

trio/_core/_multierror.py

@@ -134,11 +134,17 @@ def __exit__(self, etype, exc, tb):
             if filtered_exc is None:
                 # Swallow the exception
                 return True
-            # We can't stop Python from setting __context__, but we can
-            # hide it. (Unfortunately Python *will* wipe out any existing
-            # __context__. Nothing we can do about it :-(.)
-            filtered_exc.__suppress_context__ = True
-            raise filtered_exc
+            # When we raise filtered_exc, Python will unconditionally blow
+            # away its __context__ attribute and replace it with the original
+            # exc we caught. So after we raise it, we have to pause it while
+            # it's in flight to put the correct __context__ back.
+            old_context = filtered_exc.__context__
+            try:
+                raise filtered_exc
+            finally:
+                _, value, _ = sys.exc_info()
+                assert value is filtered_exc
+                value.__context__ = old_context
 
 class MultiError(BaseException):
     """An exception that contains other exceptions; also known as an

trio/_core/tests/test_multierror.py

@@ -228,9 +228,10 @@ def simple_filter(exc):
     # ValueError disappeared & KeyError became RuntimeError, so now:
     assert isinstance(new_m.exceptions[0], RuntimeError)
     assert isinstance(new_m.exceptions[1], NameError)
-    # we can't stop Python from attaching the original MultiError to this as a
-    # __context__, but we can hide it:
-    assert new_m.__suppress_context__
+    # Make sure that Python did not successfully attach the old MultiError to
+    # our new MultiError's __context__
+    assert not new_m.__suppress_context__
+    assert new_m.__context__ is None
 
     # check preservation of __cause__ and __context__
     v = ValueError()
@@ -241,13 +242,31 @@ def simple_filter(exc):
     assert isinstance(excinfo.value.__cause__, KeyError)
 
     v = ValueError()
-    v.__context__ = KeyError()
+    context = KeyError()
+    v.__context__ = context
     with pytest.raises(ValueError) as excinfo:
         with MultiError.catch(lambda exc: exc):
             raise v
-    assert isinstance(excinfo.value.__context__, KeyError)
+    assert excinfo.value.__context__ is context
     assert not excinfo.value.__suppress_context__
 
+    for suppress_context in [True, False]:
+        v = ValueError()
+        context = KeyError()
+        v.__context__ = context
+        v.__suppress_context__ = suppress_context
+        distractor = RuntimeError()
+        with pytest.raises(ValueError) as excinfo:
+            def catch_RuntimeError(exc):
+                if isinstance(exc, RuntimeError):
+                    return None
+                else:
+                    return exc
+            with MultiError.catch(catch_RuntimeError):
+                raise MultiError([v, distractor])
+        assert excinfo.value.__context__ is context
+        assert excinfo.value.__suppress_context__ == suppress_context
+
 
 def assert_match_in_seq(pattern_list, string):
     offset = 0