Merge #20317 from 3.3

3 files changed, 38 insertions, 1 deletions

diff --git a/Lib/contextlib.py b/Lib/contextlib.py
index d3219f6c15..ca7a79ddd1 100644
--- a/Lib/contextlib.py
+++ b/Lib/contextlib.py
@@ -298,11 +298,19 @@ class ExitStack(object):
         # we were actually nesting multiple with statements
         frame_exc = sys.exc_info()[1]
         def _fix_exception_context(new_exc, old_exc):
+            # Context isn't what we want, so find the end of the chain
             while 1:
                 exc_context = new_exc.__context__
-                if exc_context in (None, frame_exc):
+                if exc_context is old_exc:
+                    # Context is already set correctly (see issue 20317)
+                    return
+                if exc_context is None or exc_context is frame_exc:
                     break
+                details = id(new_exc), id(old_exc), id(exc_context)
+                raise Exception(str(details))
                 new_exc = exc_context
+            # Change the end of the chain to point to the exception
+            # we expect it to reference
             new_exc.__context__ = old_exc
 
         # Callbacks are invoked in LIFO order to match the behaviour of
diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py
index b8770c828f..f947232176 100644
--- a/Lib/test/test_contextlib.py
+++ b/Lib/test/test_contextlib.py
@@ -626,6 +626,29 @@ class TestExitStack(unittest.TestCase):
         else:
             self.fail("Expected KeyError, but no exception was raised")
 
+    def test_exit_exception_with_correct_context(self):
+        # http://bugs.python.org/issue20317
+        @contextmanager
+        def gets_the_context_right():
+            try:
+                yield 6
+            finally:
+                1 / 0
+
+        # The contextmanager already fixes the context, so prior to the
+        # fix, ExitStack would try to fix it *again* and get into an
+        # infinite self-referential loop
+        try:
+            with ExitStack() as stack:
+                stack.enter_context(gets_the_context_right())
+                stack.enter_context(gets_the_context_right())
+                stack.enter_context(gets_the_context_right())
+        except ZeroDivisionError as exc:
+            self.assertIsInstance(exc.__context__, ZeroDivisionError)
+            self.assertIsInstance(exc.__context__.__context__, ZeroDivisionError)
+            self.assertIsNone(exc.__context__.__context__.__context__)
+
+
     def test_body_exception_suppress(self):
         def suppress_exc(*exc_details):
             return True
diff --git a/Misc/NEWS b/Misc/NEWS
index 56c06ab36f..72ae4fa01d 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -37,6 +37,12 @@ Library
   microsecond is now rounded to one millisecond, instead of being rounded to
   zero.
 
+- Issue #20317: ExitStack.__exit__ could create a self-referential loop if an
+  exception raised by a cleanup operation already had its context set
+  correctly (for example, by the @contextmanager decorator). The infinite
+  loop this caused is now avoided by checking if the expected context is
+  already set before trying to fix it.
+
 - Issue #20311: select.epoll.poll() now rounds the timeout away from zero,
   instead of rounding towards zero. For example, a timeout of one microsecond
   is now rounded to one millisecond, instead of being rounded to zero.