pythongh-123142: Fix too wide source locations in tracebacks of exceptions from broken iterables in comprehensions

Author: iritkatriel

Lib/test/test_dictcomps.py

@@ -1,3 +1,4 @@
+import traceback
 import unittest
 
 # For scope testing.
@@ -127,6 +128,48 @@ 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
+
+        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:
+                {(x, x) for x in Iter(init_raises=True)}
+            except Exception as e:
+                return e
+
+        def next_raises():
+            try:
+                {(x, x) for x in Iter(next_raises=True)}
+            except Exception as e:
+                return e
+
+        for func, expected in [(init_raises, "Iter(init_raises=True)"),
+                               (next_raises, "Iter(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_listcomps.py

@@ -1,5 +1,6 @@
 import doctest
 import textwrap
+import traceback
 import types
 import unittest
 
@@ -711,6 +712,48 @@ 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
+
+        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:
+                [x for x in Iter(init_raises=True)]
+            except Exception as e:
+                return e
+
+        def next_raises():
+            try:
+                [x for x in Iter(next_raises=True)]
+            except Exception as e:
+                return e
+
+        for func, expected in [(init_raises, "Iter(init_raises=True)"),
+                               (next_raises, "Iter(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,4 +1,5 @@
 import doctest
+import traceback
 import unittest
 
 
@@ -148,6 +149,48 @@
 
 """
 
+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
+
+        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:
+                {x for x in Iter(init_raises=True)}
+            except Exception as e:
+                return e
+
+        def next_raises():
+            try:
+                {x for x in Iter(next_raises=True)}
+            except Exception as e:
+                return e
+
+        for func, expected in [(init_raises, "Iter(init_raises=True)"),
+                               (next_raises, "Iter(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

@@ -5174,14 +5174,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);