From a2af710c1a069749a93ca846426c6e47a20af3c0 Mon Sep 17 00:00:00 2001
From: "Nathaniel J. Smith" <njs@pobox.com>
Date: Thu, 18 May 2017 22:15:05 -0700
Subject: [PATCH] Correctly preserve exception __context__ in MultiError.catch

Python's implicit exception chaining logic insists on corrupting
__context__ when we re-raise an unpacked exception in
MultiError.catch. This commit introduces a counter-measure.
---
 trio/_core/_multierror.py           | 16 +++++++++++-----
 trio/_core/tests/test_multierror.py | 29 ++++++++++++++++++++++++-----
 2 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/trio/_core/_multierror.py b/trio/_core/_multierror.py
index 67ce82b176..56c9b404f9 100644
--- a/trio/_core/_multierror.py
+++ b/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
diff --git a/trio/_core/tests/test_multierror.py b/trio/_core/tests/test_multierror.py
index 37e01b7d3b..14989ee386 100644
--- a/trio/_core/tests/test_multierror.py
+++ b/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