gh-123142: Fix too wide source locations in tracebacks of exceptions from broken iterables in comprehensions (#123173)

iritkatriel · web-flow · commit ec89620e5e14 · 2024-08-21T19:12:05.000+01:00
diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py
@@ -59,7 +59,8 @@
     "Py_DEBUG", "exceeds_recursion_limit", "get_c_recursion_limit",
     "skip_on_s390x",
     "without_optimizer",
-    "force_not_colorized"
+    "force_not_colorized",
+    "BrokenIter",
     ]
 
 
@@ -2847,3 +2848,16 @@ def get_signal_name(exitcode):
         pass
 
     return None
+
+class BrokenIter:
+    def __init__(self, init_raises=False, next_raises=False):
+        if init_raises:
+            1/0
+        self.next_raises = next_raises
+
+    def __next__(self):
+        if self.next_raises:
+            1/0
+
+    def __iter__(self):
+        return self
diff --git a/Lib/test/test_compile.py b/Lib/test/test_compile.py
@@ -1172,7 +1172,7 @@ def return_genexp():
                     x
                     in
                     y)
-        genexp_lines = [0, 2, 0]
+        genexp_lines = [0, 4, 2, 0, 4]
 
         genexp_code = return_genexp.__code__.co_consts[1]
         code_lines = self.get_code_lines(genexp_code)
@@ -1627,7 +1627,7 @@ def test_multiline_generator_expression(self):
         self.assertOpcodeSourcePositionIs(compiled_code, 'JUMP_BACKWARD',
             line=1, end_line=2, column=1, end_column=8, occurrence=1)
         self.assertOpcodeSourcePositionIs(compiled_code, 'RETURN_CONST',
-            line=1, end_line=6, column=0, end_column=32, occurrence=1)
+            line=4, end_line=4, column=7, end_column=14, occurrence=1)
 
     def test_multiline_async_generator_expression(self):
         snippet = textwrap.dedent("""\
diff --git a/Lib/test/test_dictcomps.py b/Lib/test/test_dictcomps.py
@@ -1,5 +1,8 @@
+import traceback
 import unittest
 
+from test.support import BrokenIter
+
 # For scope testing.
 g = "Global variable"
 
@@ -127,6 +130,34 @@ def test_star_expression(self):
         self.assertEqual({i: i*i for i in [*range(4)]}, expected)
         self.assertEqual({i: i*i for i in (*range(4),)}, expected)
 
+    def test_exception_locations(self):
+        # The location of an exception raised from __init__ or
+        # __next__ should should be the iterator expression
+        def init_raises():
+            try:
+                {x:x for x in BrokenIter(init_raises=True)}
+            except Exception as e:
+                return e
+
+        def next_raises():
+            try:
+                {x:x for x in BrokenIter(next_raises=True)}
+            except Exception as e:
+                return e
+
+        for func, expected in [(init_raises, "BrokenIter(init_raises=True)"),
+                               (next_raises, "BrokenIter(next_raises=True)"),
+                              ]:
+            with self.subTest(func):
+                exc = func()
+                f = traceback.extract_tb(exc.__traceback__)[0]
+                indent = 16
+                co = func.__code__
+                self.assertEqual(f.lineno, co.co_firstlineno + 2)
+                self.assertEqual(f.end_lineno, co.co_firstlineno + 2)
+                self.assertEqual(f.line[f.colno - indent : f.end_colno - indent],
+                                 expected)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Lib/test/test_iter.py b/Lib/test/test_iter.py
@@ -5,6 +5,7 @@
 from test.support import cpython_only
 from test.support.os_helper import TESTFN, unlink
 from test.support import check_free_after_iterating, ALWAYS_EQ, NEVER_EQ
+from test.support import BrokenIter
 import pickle
 import collections.abc
 import functools
@@ -1148,35 +1149,22 @@ def test_exception_locations(self):
         # The location of an exception raised from __init__ or
         # __next__ should should be the iterator expression
 
-        class Iter:
-            def __init__(self, init_raises=False, next_raises=False):
-                if init_raises:
-                    1/0
-                self.next_raises = next_raises
-
-            def __next__(self):
-                if self.next_raises:
-                    1/0
-
-            def __iter__(self):
-                return self
-
         def init_raises():
             try:
-                for x in Iter(init_raises=True):
+                for x in BrokenIter(init_raises=True):
                     pass
             except Exception as e:
                 return e
 
         def next_raises():
             try:
-                for x in Iter(next_raises=True):
+                for x in BrokenIter(next_raises=True):
                     pass
             except Exception as e:
                 return e
 
-        for func, expected in [(init_raises, "Iter(init_raises=True)"),
-                               (next_raises, "Iter(next_raises=True)"),
+        for func, expected in [(init_raises, "BrokenIter(init_raises=True)"),
+                               (next_raises, "BrokenIter(next_raises=True)"),
                               ]:
             with self.subTest(func):
                 exc = func()
diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py
@@ -1,8 +1,11 @@
 import doctest
 import textwrap
+import traceback
 import types
 import unittest
 
+from test.support import BrokenIter
+
 
 doctests = """
 ########### Tests borrowed from or inspired by test_genexps.py ############
@@ -711,6 +714,35 @@ def test_multiple_comprehension_name_reuse(self):
         self._check_in_scopes(code, {"x": 2, "y": [3]}, ns={"x": 3}, scopes=["class"])
         self._check_in_scopes(code, {"x": 2, "y": [2]}, ns={"x": 3}, scopes=["function", "module"])
 
+    def test_exception_locations(self):
+        # The location of an exception raised from __init__ or
+        # __next__ should should be the iterator expression
+
+        def init_raises():
+            try:
+                [x for x in BrokenIter(init_raises=True)]
+            except Exception as e:
+                return e
+
+        def next_raises():
+            try:
+                [x for x in BrokenIter(next_raises=True)]
+            except Exception as e:
+                return e
+
+        for func, expected in [(init_raises, "BrokenIter(init_raises=True)"),
+                               (next_raises, "BrokenIter(next_raises=True)"),
+                              ]:
+            with self.subTest(func):
+                exc = func()
+                f = traceback.extract_tb(exc.__traceback__)[0]
+                indent = 16
+                co = func.__code__
+                self.assertEqual(f.lineno, co.co_firstlineno + 2)
+                self.assertEqual(f.end_lineno, co.co_firstlineno + 2)
+                self.assertEqual(f.line[f.colno - indent : f.end_colno - indent],
+                                 expected)
+
 __test__ = {'doctests' : doctests}
 
 def load_tests(loader, tests, pattern):
diff --git a/Lib/test/test_setcomps.py b/Lib/test/test_setcomps.py
@@ -1,6 +1,9 @@
 import doctest
+import traceback
 import unittest
 
+from test.support import BrokenIter
+
 
 doctests = """
 ########### Tests mostly copied from test_listcomps.py ############
@@ -148,6 +151,35 @@
 
 """
 
+class SetComprehensionTest(unittest.TestCase):
+    def test_exception_locations(self):
+        # The location of an exception raised from __init__ or
+        # __next__ should should be the iterator expression
+
+        def init_raises():
+            try:
+                {x for x in BrokenIter(init_raises=True)}
+            except Exception as e:
+                return e
+
+        def next_raises():
+            try:
+                {x for x in BrokenIter(next_raises=True)}
+            except Exception as e:
+                return e
+
+        for func, expected in [(init_raises, "BrokenIter(init_raises=True)"),
+                               (next_raises, "BrokenIter(next_raises=True)"),
+                              ]:
+            with self.subTest(func):
+                exc = func()
+                f = traceback.extract_tb(exc.__traceback__)[0]
+                indent = 16
+                co = func.__code__
+                self.assertEqual(f.lineno, co.co_firstlineno + 2)
+                self.assertEqual(f.end_lineno, co.co_firstlineno + 2)
+                self.assertEqual(f.line[f.colno - indent : f.end_colno - indent],
+                                 expected)
 
 __test__ = {'doctests' : doctests}
 
diff --git a/Misc/NEWS.d/next/Core and Builtins/2024-08-20-12-29-52.gh-issue-123142.3PXiNb.rst b/Misc/NEWS.d/next/Core and Builtins/2024-08-20-12-29-52.gh-issue-123142.3PXiNb.rst
@@ -0,0 +1,2 @@
+Fix too-wide source location in exception tracebacks coming from broken
+iterables in comprehensions.
diff --git a/Python/compile.c b/Python/compile.c
@@ -5183,14 +5183,15 @@ compiler_sync_comprehension_generator(struct compiler *c, location loc,
             }
             if (IS_LABEL(start)) {
                 VISIT(c, expr, gen->iter);
-                ADDOP(c, loc, GET_ITER);
+                ADDOP(c, LOC(gen->iter), GET_ITER);
             }
         }
     }
+
     if (IS_LABEL(start)) {
         depth++;
         USE_LABEL(c, start);
-        ADDOP_JUMP(c, loc, FOR_ITER, anchor);
+        ADDOP_JUMP(c, LOC(gen->iter), FOR_ITER, anchor);
     }
     VISIT(c, expr, gen->target);
 

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+Fix too-wide source location in exception tracebacks coming from broken`
	`2`	`+iterables in comprehensions.`
Original file line number	Diff line number	Diff line change
`@@ -5183,14 +5183,15 @@ compiler_sync_comprehension_generator(struct compiler *c, location loc,`
`5183`	`5183`	`}`
`5184`	`5184`	`if (IS_LABEL(start)) {`
`5185`	`5185`	`VISIT(c, expr, gen->iter);`
`5186`		`- ADDOP(c, loc, GET_ITER);`
	`5186`	`+ ADDOP(c, LOC(gen->iter), GET_ITER);`
`5187`	`5187`	`}`
`5188`	`5188`	`}`
`5189`	`5189`	`}`
	`5190`	`+`
`5190`	`5191`	`if (IS_LABEL(start)) {`
`5191`	`5192`	`depth++;`
`5192`	`5193`	`USE_LABEL(c, start);`
`5193`		`- ADDOP_JUMP(c, loc, FOR_ITER, anchor);`
	`5194`	`+ ADDOP_JUMP(c, LOC(gen->iter), FOR_ITER, anchor);`
`5194`	`5195`	`}`
`5195`	`5196`	`VISIT(c, expr, gen->target);`
`5196`	`5197`