From 233f6b644fe7fbe0aeffb7005f56c65385400de5 Mon Sep 17 00:00:00 2001
From: Mikhail Efimov <efimov.mikhail@gmail.com>
Date: Tue, 15 Oct 2024 13:14:22 +0300
Subject: [PATCH] gh-125038: Changes in tests on generator modifying

New tests are moved back to test_generators.py.
Tests on generator creation via FunctionType from gi_code are added.
---
 Lib/test/test_frame.py      | 55 --------------------------
 Lib/test/test_generators.py | 78 +++++++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 55 deletions(-)

diff --git a/Lib/test/test_frame.py b/Lib/test/test_frame.py
index 341b64280f9ce1..32de8ed9a13f80 100644
--- a/Lib/test/test_frame.py
+++ b/Lib/test/test_frame.py
@@ -515,61 +515,6 @@ def make_frame():
         with self.assertRaises(TypeError):
             FrameLocalsProxy(frame=sys._getframe())  # no keyword arguments
 
-    def test_generator_f_locals(self):
-        def get_generator_genexpr(new_value):
-            g = (x for x in range(10))
-            g.gi_frame.f_locals['.0'] = new_value
-            return g
-
-        def get_generator_fn_call(new_value):
-            def gen(it):
-                for x in it:
-                    yield x
-
-            g = gen(range(10))
-            g.gi_frame.f_locals['it'] = new_value
-            return g
-
-        err_msg_pattern_genexpr = "'%s' object is not an iterator"
-        err_msg_pattern_fn_call = "'%s' object is not iterable"
-
-        sequences = [
-            range(0),
-            range(20),
-            [1, 2, 3],
-            (2,),
-            set((13, 48, 211)),
-            frozenset((15, 8, 6)),
-            dict([(1, 2), (3, 4)]),
-        ]
-
-        for seq in sequences:
-            err_msg_genexpr = err_msg_pattern_genexpr % type(seq).__name__
-            with self.assertRaisesRegex(TypeError, err_msg_genexpr):
-                list(get_generator_genexpr(seq))
-            self.assertListEqual(list(get_generator_genexpr(iter(seq))),
-                                 list(seq))
-
-            self.assertListEqual(list(get_generator_fn_call(seq)),
-                                 list(seq))
-            self.assertListEqual(list(get_generator_fn_call(iter(seq))),
-                                 list(seq))
-
-        non_sequences = [
-            None,
-            42,
-            3.0,
-            2j,
-        ]
-
-        for obj in non_sequences:
-            err_msg_genexpr = err_msg_pattern_genexpr % type(obj).__name__
-            with self.assertRaisesRegex(TypeError, err_msg_genexpr):
-                list(get_generator_genexpr(obj))
-            err_msg_fn_call = err_msg_pattern_fn_call % type(obj).__name__
-            with self.assertRaisesRegex(TypeError, err_msg_fn_call):
-                list(get_generator_fn_call(obj))
-
 
 class FrameLocalsProxyMappingTests(mapping_tests.TestHashMappingProtocol):
     """Test that FrameLocalsProxy behaves like a Mapping (with exceptions)"""
diff --git a/Lib/test/test_generators.py b/Lib/test/test_generators.py
index 03a31ec6a05726..58a1e73c11ef1a 100644
--- a/Lib/test/test_generators.py
+++ b/Lib/test/test_generators.py
@@ -268,6 +268,84 @@ def loop():
         #This should not raise
         loop()
 
+
+class ModifyTest(unittest.TestCase):
+    sequences = [
+        range(0),
+        range(20),
+        [1, 2, 3],
+        (2,),
+        set((13, 48, 211)),
+        frozenset((15, 8, 6)),
+        dict([(1, 2), (3, 4)]),
+    ]
+
+    non_sequences = [
+        None,
+        42,
+        3.0,
+        2j,
+    ]
+
+    def modify_f_locals(self, g, local, new_iter):
+        g.gi_frame.f_locals[local] = new_iter
+        return g
+
+    def new_gen_from_gi_code(self, g, new_iter):
+        generator_func = types.FunctionType(g.gi_code, {})
+        return generator_func(new_iter)
+
+    def test_modify_genexpr(self):
+        def genexpr():
+            return (x for x in range(10))
+
+        def get_genexpr_f_locals(self, new_iter):
+            return self.modify_f_locals(genexpr(), '.0', new_iter)
+
+        def get_genexpr_gi_code(self, new_iter):
+            return self.new_gen_from_gi_code(genexpr(), new_iter)
+
+        err_msg_pattern_genexpr = "'%s' object is not an iterator"
+
+        for get_genexpr in [get_genexpr_f_locals, get_genexpr_gi_code]:
+            for seq in self.sequences:
+                err_msg_genexpr = err_msg_pattern_genexpr % type(seq).__name__
+                with self.assertRaisesRegex(TypeError, err_msg_genexpr):
+                    list(get_genexpr(self, seq))
+                self.assertListEqual(list(get_genexpr(self, iter(seq))),
+                                     list(seq))
+            for obj in self.non_sequences:
+                err_msg_genexpr = err_msg_pattern_genexpr % type(obj).__name__
+                with self.assertRaisesRegex(TypeError, err_msg_genexpr):
+                    list(get_genexpr(self, obj))
+
+    def test_modify_genfunc(self):
+        def genfunc():
+            def gen(it):
+                for x in it:
+                    yield x
+            return gen(range(10))
+
+        def get_genfunc_f_locals(self, new_iter):
+            return self.modify_f_locals(genfunc(), 'it', new_iter)
+
+        def get_genfunc_gi_code(self, new_iter):
+            return self.new_gen_from_gi_code(genfunc(), new_iter)
+
+        err_msg_pattern_fn_call = "'%s' object is not iterable"
+
+        for get_genfunc in [get_genfunc_f_locals, get_genfunc_gi_code]:
+            for seq in self.sequences:
+                self.assertListEqual(list(get_genfunc(self, seq)),
+                                     list(seq))
+                self.assertListEqual(list(get_genfunc(self, iter(seq))),
+                                     list(seq))
+            for obj in self.non_sequences:
+                err_msg_fn_call = err_msg_pattern_fn_call % type(obj).__name__
+                with self.assertRaisesRegex(TypeError, err_msg_fn_call):
+                    list(get_genfunc(self, obj))
+
+
 class ExceptionTest(unittest.TestCase):
     # Tests for the issue #23353: check that the currently handled exception
     # is correctly saved/restored in PyEval_EvalFrameEx().