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

Lib/test/support/__init__.py

@@ -59,7 +59,8 @@ __all__ = [
     "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

Lib/test/test_compile.py

@@ -1172,7 +1172,7 @@ class TestSpecifics(unittest.TestCase):
                     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 @@ class TestSourcePositions(unittest.TestCase):
         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("""\

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 @@ class DictComprehensionTest(unittest.TestCase):
         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()

Lib/test/test_iter.py

@@ -5,6 +5,7 @@ import unittest
 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 @@ class TestCase(unittest.TestCase):
         # 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)"),
+        for func, expected in [(init_raises, "BrokenIter(init_raises=True)"),
-                               (next_raises, "Iter(next_raises=True)"),
+                               (next_raises, "BrokenIter(next_raises=True)"),
                               ]:
             with self.subTest(func):
                 exc = func()

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 @@ class ListComprehensionTest(unittest.TestCase):
         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):

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 @@ We also repeat each of the above scoping tests inside a function
 
 """
 
+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}
 

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.

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);