From add94c9d825dc2d870f6557e1d6510a82d3dba12 Mon Sep 17 00:00:00 2001
From: Nick Coghlan <ncoghlan@gmail.com>
Date: Fri, 24 Jan 2014 23:05:45 +1000
Subject: [PATCH] Issue 20317: Remove debugging code from contextlib

- Alex J Burke noticed a debugging raise in the commit that
  fixed the original bug reported in issue 20317
- this showed that multiple iterations through the affected
  loop wasn't actually being tested
---
 Lib/contextlib.py           |  4 +--
 Lib/test/test_contextlib.py | 58 ++++++++++++++++++++++++++++++-------
 2 files changed, 49 insertions(+), 13 deletions(-)

diff --git a/Lib/contextlib.py b/Lib/contextlib.py
index f87828569a1..b03f8283194 100644
--- a/Lib/contextlib.py
+++ b/Lib/contextlib.py
@@ -231,7 +231,7 @@ 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
+            # Context may not be correct, so find the end of the chain
             while 1:
                 exc_context = new_exc.__context__
                 if exc_context is old_exc:
@@ -239,8 +239,6 @@ class ExitStack(object):
                     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
diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py
index e5365f78259..e38d043a5bc 100644
--- a/Lib/test/test_contextlib.py
+++ b/Lib/test/test_contextlib.py
@@ -603,24 +603,62 @@ class TestExitStack(unittest.TestCase):
     def test_exit_exception_with_correct_context(self):
         # http://bugs.python.org/issue20317
         @contextmanager
-        def gets_the_context_right():
+        def gets_the_context_right(exc):
             try:
-                yield 6
+                yield
             finally:
-                1 / 0
+                raise exc
+
+        exc1 = Exception(1)
+        exc2 = Exception(2)
+        exc3 = Exception(3)
+        exc4 = Exception(4)
 
         # 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__)
+                stack.enter_context(gets_the_context_right(exc4))
+                stack.enter_context(gets_the_context_right(exc3))
+                stack.enter_context(gets_the_context_right(exc2))
+                raise exc1
+        except Exception as exc:
+            self.assertIs(exc, exc4)
+            self.assertIs(exc.__context__, exc3)
+            self.assertIs(exc.__context__.__context__, exc2)
+            self.assertIs(exc.__context__.__context__.__context__, exc1)
+            self.assertIsNone(
+                       exc.__context__.__context__.__context__.__context__)
+
+    def test_exit_exception_with_existing_context(self):
+        # Addresses a lack of test coverage discovered after checking in a
+        # fix for issue 20317 that still contained debugging code.
+        def raise_nested(inner_exc, outer_exc):
+            try:
+                raise inner_exc
+            finally:
+                raise outer_exc
+        exc1 = Exception(1)
+        exc2 = Exception(2)
+        exc3 = Exception(3)
+        exc4 = Exception(4)
+        exc5 = Exception(5)
+        try:
+            with ExitStack() as stack:
+                stack.callback(raise_nested, exc4, exc5)
+                stack.callback(raise_nested, exc2, exc3)
+                raise exc1
+        except Exception as exc:
+            self.assertIs(exc, exc5)
+            self.assertIs(exc.__context__, exc4)
+            self.assertIs(exc.__context__.__context__, exc3)
+            self.assertIs(exc.__context__.__context__.__context__, exc2)
+            self.assertIs(
+                 exc.__context__.__context__.__context__.__context__, exc1)
+            self.assertIsNone(
+                exc.__context__.__context__.__context__.__context__.__context__)
+
 
 
     def test_body_exception_suppress(self):