gh-110543: Fix CodeType.replace in presence of comprehensions (#110586)

1 files changed, 43 insertions, 2 deletions

diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py
index 12f7bbd..f95a78a 100644
--- a/Lib/test/test_listcomps.py
+++ b/Lib/test/test_listcomps.py
@@ -1,5 +1,6 @@
 import doctest
 import textwrap
+import types
 import unittest
 
 
@@ -92,7 +93,8 @@ Make sure that None is a valid return value
 
 
 class ListComprehensionTest(unittest.TestCase):
-    def _check_in_scopes(self, code, outputs=None, ns=None, scopes=None, raises=()):
+    def _check_in_scopes(self, code, outputs=None, ns=None, scopes=None, raises=(),
+                         exec_func=exec):
         code = textwrap.dedent(code)
         scopes = scopes or ["module", "class", "function"]
         for scope in scopes:
@@ -119,7 +121,7 @@ class ListComprehensionTest(unittest.TestCase):
                         return moddict[name]
                 newns = ns.copy() if ns else {}
                 try:
-                    exec(newcode, newns)
+                    exec_func(newcode, newns)
                 except raises as e:
                     # We care about e.g. NameError vs UnboundLocalError
                     self.assertIs(type(e), raises)
@@ -613,6 +615,45 @@ class ListComprehensionTest(unittest.TestCase):
         import sys
         self._check_in_scopes(code, {"val": 0}, ns={"sys": sys})
 
+    def _recursive_replace(self, maybe_code):
+        if not isinstance(maybe_code, types.CodeType):
+            return maybe_code
+        return maybe_code.replace(co_consts=tuple(
+            self._recursive_replace(c) for c in maybe_code.co_consts
+        ))
+
+    def _replacing_exec(self, code_string, ns):
+        co = compile(code_string, "<string>", "exec")
+        co = self._recursive_replace(co)
+        exec(co, ns)
+
+    def test_code_replace(self):
+        code = """
+            x = 3
+            [x for x in (1, 2)]
+            dir()
+            y = [x]
+        """
+        self._check_in_scopes(code, {"y": [3], "x": 3})
+        self._check_in_scopes(code, {"y": [3], "x": 3}, exec_func=self._replacing_exec)
+
+    def test_code_replace_extended_arg(self):
+        num_names = 300
+        assignments = "; ".join(f"x{i} = {i}" for i in range(num_names))
+        name_list = ", ".join(f"x{i}" for i in range(num_names))
+        expected = {
+            "y": list(range(num_names)),
+            **{f"x{i}": i for i in range(num_names)}
+        }
+        code = f"""
+            {assignments}
+            [({name_list}) for {name_list} in (range(300),)]
+            dir()
+            y = [{name_list}]
+        """
+        self._check_in_scopes(code, expected)
+        self._check_in_scopes(code, expected, exec_func=self._replacing_exec)
+
 
 __test__ = {'doctests' : doctests}