Make sure that line number is set correctly for call to __exit__ when handling exception in body of a with statement. (GH-26890)

1 files changed, 24 insertions, 7 deletions

diff --git a/Lib/test/test_exceptions.py b/Lib/test/test_exceptions.py
index d444a12..8f68954 100644
--- a/Lib/test/test_exceptions.py
+++ b/Lib/test/test_exceptions.py
@@ -2160,18 +2160,23 @@ class SyntaxErrorTests(unittest.TestCase):
 
 class PEP626Tests(unittest.TestCase):
 
-    def lineno_after_raise(self, f, line):
+    def lineno_after_raise(self, f, *expected):
         try:
             f()
         except Exception as ex:
             t = ex.__traceback__
-            while t.tb_next:
-                t = t.tb_next
+        else:
+            self.fail("No exception raised")
+        lines = []
+        t = t.tb_next # Skip this function
+        while t:
             frame = t.tb_frame
-            if line is None:
-                self.assertEqual(frame.f_lineno, line)
-            else:
-                self.assertEqual(frame.f_lineno-frame.f_code.co_firstlineno, line)
+            lines.append(
+                None if frame.f_lineno is None else
+                frame.f_lineno-frame.f_code.co_firstlineno
+            )
+            t = t.tb_next
+        self.assertEqual(tuple(lines), expected)
 
     def test_lineno_after_raise_simple(self):
         def simple():
@@ -2250,5 +2255,17 @@ class PEP626Tests(unittest.TestCase):
         f.__code__ = f.__code__.replace(co_linetable=b'\x04\x80\xff\x80')
         self.lineno_after_raise(f, None)
 
+    def test_lineno_after_raise_in_with_exit(self):
+        class ExitFails:
+            def __enter__(self):
+                return self
+            def __exit__(self, *args):
+                raise ValueError
+
+        def after_with():
+            with ExitFails():
+                1/0
+        self.lineno_after_raise(after_with, 1, 1)
+
 if __name__ == '__main__':
     unittest.main()