gh-125038: Fix crash after genexpr.gi_frame.f_locals manipulations (#125178)

efimov-mikhail · web-flow · commit 079875e39589 · 2024-10-22T09:00:25.000-07:00
diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
@@ -810,6 +810,7 @@ def foo(x):
                POP_TOP
        L1:     RESUME                   0
                LOAD_FAST                0 (.0)
+               GET_ITER
        L2:     FOR_ITER                10 (to L3)
                STORE_FAST               1 (z)
                LOAD_DEREF               2 (x)
diff --git a/Lib/test/test_generators.py b/Lib/test/test_generators.py
@@ -268,6 +268,79 @@ def loop():
         #This should not raise
         loop()
 
+
+class ModifyUnderlyingIterableTest(unittest.TestCase):
+    iterables = [
+        range(0),
+        range(20),
+        [1, 2, 3],
+        (2,),
+        {13, 48, 211},
+        frozenset((15, 8, 6)),
+        {1: 2, 3: 4},
+    ]
+
+    non_iterables = [
+        None,
+        42,
+        3.0,
+        2j,
+    ]
+
+    def genexpr(self):
+        return (x for x in range(10))
+
+    def genfunc(self):
+        def gen(it):
+            for x in it:
+                yield x
+        return gen(range(10))
+
+    def process_tests(self, get_generator):
+        for obj in self.iterables:
+            g_obj = get_generator(obj)
+            with self.subTest(g_obj=g_obj, obj=obj):
+                self.assertListEqual(list(g_obj), list(obj))
+
+            g_iter = get_generator(iter(obj))
+            with self.subTest(g_iter=g_iter, obj=obj):
+                self.assertListEqual(list(g_iter), list(obj))
+
+        err_regex = "'.*' object is not iterable"
+        for obj in self.non_iterables:
+            g_obj = get_generator(obj)
+            with self.subTest(g_obj=g_obj):
+                self.assertRaisesRegex(TypeError, err_regex, list, g_obj)
+
+    def test_modify_f_locals(self):
+        def modify_f_locals(g, local, obj):
+            g.gi_frame.f_locals[local] = obj
+            return g
+
+        def get_generator_genexpr(obj):
+            return modify_f_locals(self.genexpr(), '.0', obj)
+
+        def get_generator_genfunc(obj):
+            return modify_f_locals(self.genfunc(), 'it', obj)
+
+        self.process_tests(get_generator_genexpr)
+        self.process_tests(get_generator_genfunc)
+
+    def test_new_gen_from_gi_code(self):
+        def new_gen_from_gi_code(g, obj):
+            generator_func = types.FunctionType(g.gi_code, {})
+            return generator_func(obj)
+
+        def get_generator_genexpr(obj):
+            return new_gen_from_gi_code(self.genexpr(), obj)
+
+        def get_generator_genfunc(obj):
+            return new_gen_from_gi_code(self.genfunc(), obj)
+
+        self.process_tests(get_generator_genexpr)
+        self.process_tests(get_generator_genfunc)
+
+
 class ExceptionTest(unittest.TestCase):
     # Tests for the issue #23353: check that the currently handled exception
     # is correctly saved/restored in PyEval_EvalFrameEx().
diff --git a/Misc/NEWS.d/next/Core_and_Builtins/2024-10-09-13-53-50.gh-issue-125038.ffSLCz.rst b/Misc/NEWS.d/next/Core_and_Builtins/2024-10-09-13-53-50.gh-issue-125038.ffSLCz.rst
@@ -0,0 +1,2 @@
+Fix crash when iterating over a generator expression after direct changes on ``gi_frame.f_locals``.
+Patch by Mikhail Efimov.
diff --git a/Python/codegen.c b/Python/codegen.c
@@ -4164,6 +4164,7 @@ codegen_sync_comprehension_generator(compiler *c, location loc,
 
     if (IS_JUMP_TARGET_LABEL(start)) {
         depth++;
+        ADDOP(c, LOC(gen->iter), GET_ITER);
         USE_LABEL(c, start);
         ADDOP_JUMP(c, LOC(gen->iter), FOR_ITER, anchor);
     }

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	+Fix crash when iterating over a generator expression after direct changes on ``gi_frame.f_locals``.
	`2`	`+Patch by Mikhail Efimov.`
Original file line number	Diff line number	Diff line change
`@@ -4164,6 +4164,7 @@ codegen_sync_comprehension_generator(compiler *c, location loc,`
`4164`	`4164`
`4165`	`4165`	`if (IS_JUMP_TARGET_LABEL(start)) {`
`4166`	`4166`	`depth++;`
	`4167`	`+ ADDOP(c, LOC(gen->iter), GET_ITER);`
`4167`	`4168`	`USE_LABEL(c, start);`
`4168`	`4169`	`ADDOP_JUMP(c, LOC(gen->iter), FOR_ITER, anchor);`
`4169`	`4170`	`}`