From 079875e39589eb0628b5883f7ffa387e7476ec06 Mon Sep 17 00:00:00 2001
From: Mikhail Efimov <efimov.mikhail@gmail.com>
Date: Tue, 22 Oct 2024 19:00:25 +0300
Subject: [PATCH] gh-125038: Fix crash after genexpr.gi_frame.f_locals
 manipulations (#125178)

---
 Lib/test/test_dis.py                          |  1 +
 Lib/test/test_generators.py                   | 73 +++++++++++++++++++
 ...-10-09-13-53-50.gh-issue-125038.ffSLCz.rst |  2 +
 Python/codegen.c                              |  1 +
 4 files changed, 77 insertions(+)
 create mode 100644 Misc/NEWS.d/next/Core_and_Builtins/2024-10-09-13-53-50.gh-issue-125038.ffSLCz.rst

diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
index 1ee0fbe9891..1f9c04cdbc9 100644
--- a/Lib/test/test_dis.py
+++ b/Lib/test/test_dis.py
@@ -810,6 +810,7 @@ Disassembly of <code object <genexpr> at 0x..., file "%s", line %d>:
                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
index 03a31ec6a05..bf2cb116072 100644
--- a/Lib/test/test_generators.py
+++ b/Lib/test/test_generators.py
@@ -268,6 +268,79 @@ class GeneratorTest(unittest.TestCase):
         #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
new file mode 100644
index 00000000000..15de48ec0e4
--- /dev/null
+++ 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
index 689d2b5124e..bfacc6f0c55 100644
--- 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);
     }