[3.12] gh-93691: fix too broad source locations of for statement iterators (GH-120330 (#120405)

[3.12] gh-93691: fix too broad source locations of for statement iterators (GH-120330).
(cherry picked from commit 97b69db167)

Lib/test/test_compiler_codegen.py

@@ -39,6 +39,7 @@ class IsolatedCodeGenTests(CodegenTestCase):
             ('GET_ITER', None, 1),
             loop_lbl := self.Label(),
             ('FOR_ITER', exit_lbl := self.Label(), 1),
+            ('NOP', None, 1, 1),
             ('STORE_NAME', 1, 1),
             ('PUSH_NULL', None, 2),
             ('LOAD_NAME', 2, 2),

Lib/test/test_iter.py

@@ -10,6 +10,7 @@ import collections.abc
 import functools
 import contextlib
 import builtins
+import traceback
 
 # Test result of triple loop (too big to inline)
 TRIPLETS = [(0, 0, 0), (0, 0, 1), (0, 0, 2),
@@ -1143,6 +1144,51 @@ class TestCase(unittest.TestCase):
             self.assertRaises(TypeError, iter, typ())
         self.assertRaises(ZeroDivisionError, iter, BadIterableClass())
 
+    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):
+                    pass
+            except Exception as e:
+                return e
+
+        def next_raises():
+            try:
+                for x in Iter(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)"),
+                              ]:
+            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_sys_settrace.py

@@ -1634,15 +1634,15 @@ class TraceTestCase(unittest.TestCase):
         EXPECTED_EVENTS = [
             (0, 'call'),
             (2, 'line'),
-            (1, 'line'),
             (-3, 'call'),
             (-2, 'line'),
             (-2, 'return'),
-            (4, 'line'),
             (1, 'line'),
+            (4, 'line'),
+            (2, 'line'),
             (-2, 'call'),
             (-2, 'return'),
-            (1, 'return'),
+            (2, 'return'),
         ]
 
         # C level events should be the same as expected and the same as Python level.

Misc/NEWS.d/next/Core and Builtins/2024-06-10-22-30-26.gh-issue-93691.68WOTS.rst

@@ -0,0 +1,2 @@
+Fix source locations of instructions generated for the iterator of a for
+statement.

Programs/test_frozenmain.h

@@ -27,12 +27,11 @@ unsigned char M_test_frozenmain[] = {
     218,3,107,101,121,169,0,243,0,0,0,0,250,18,116,101,
     115,116,95,102,114,111,122,101,110,109,97,105,110,46,112,121,
     250,8,60,109,111,100,117,108,101,62,114,18,0,0,0,1,
-    0,0,0,115,102,0,0,0,240,3,1,1,1,243,8,0,
+    0,0,0,115,97,0,0,0,240,3,1,1,1,243,8,0,
     1,11,219,0,24,225,0,5,208,6,26,212,0,27,217,0,
     5,128,106,144,35,151,40,145,40,212,0,27,216,9,38,208,
     9,26,215,9,38,209,9,38,211,9,40,168,24,209,9,50,
-    128,6,240,2,6,12,2,242,0,7,1,42,128,67,241,14,
-    0,5,10,136,71,144,67,144,53,152,2,152,54,160,35,153,
-    59,152,45,208,10,40,213,4,41,241,15,7,1,42,114,16,
-    0,0,0,
+    128,6,243,2,6,12,2,128,67,241,14,0,5,10,136,71,
+    144,67,144,53,152,2,152,54,160,35,153,59,152,45,208,10,
+    40,213,4,41,241,15,6,12,2,114,16,0,0,0,
 };

Python/compile.c

@@ -3024,11 +3024,18 @@ compiler_for(struct compiler *c, stmt_ty s)
     RETURN_IF_ERROR(compiler_push_fblock(c, loc, FOR_LOOP, start, end, NULL));
 
     VISIT(c, expr, s->v.For.iter);
+
+    loc = LOC(s->v.For.iter);
     ADDOP(c, loc, GET_ITER);
 
     USE_LABEL(c, start);
     ADDOP_JUMP(c, loc, FOR_ITER, cleanup);
 
+    /* Add NOP to ensure correct line tracing of multiline for statements.
+     * It will be removed later if redundant.
+     */
+    ADDOP(c, LOC(s->v.For.target), NOP);
+
     USE_LABEL(c, body);
     VISIT(c, expr, s->v.For.target);
     VISIT_SEQ(c, stmt, s->v.For.body);