From 0b718e6407da65b838576a2459d630824ca62155 Mon Sep 17 00:00:00 2001
From: Jelle Zijlstra <jelle.zijlstra@gmail.com>
Date: Wed, 8 Nov 2023 12:11:59 -0800
Subject: [PATCH] gh-110543: Fix CodeType.replace in presence of comprehensions
 (#110586)

---
 Lib/test/test_listcomps.py                    | 45 ++++++++++++++++++-
 ...-10-09-19-54-33.gh-issue-110543.1wrxO8.rst |  3 ++
 Objects/codeobject.c                          | 29 ++++++++++++
 3 files changed, 75 insertions(+), 2 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Core and Builtins/2023-10-09-19-54-33.gh-issue-110543.1wrxO8.rst

diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py
index 12f7bbd123b..f95a78aff0c 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}
 
diff --git a/Misc/NEWS.d/next/Core and Builtins/2023-10-09-19-54-33.gh-issue-110543.1wrxO8.rst b/Misc/NEWS.d/next/Core and Builtins/2023-10-09-19-54-33.gh-issue-110543.1wrxO8.rst
new file mode 100644
index 00000000000..5f9571566da
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2023-10-09-19-54-33.gh-issue-110543.1wrxO8.rst	
@@ -0,0 +1,3 @@
+Fix regression in Python 3.12 where :meth:`types.CodeType.replace` would
+produce a broken code object if called on a module or class code object that
+contains a comprehension. Patch by Jelle Zijlstra.
diff --git a/Objects/codeobject.c b/Objects/codeobject.c
index 79ac5749037..dc46b773c26 100644
--- a/Objects/codeobject.c
+++ b/Objects/codeobject.c
@@ -643,6 +643,35 @@ PyUnstable_Code_NewWithPosOnlyArgs(
         _Py_set_localsplus_info(offset, name, CO_FAST_FREE,
                                localsplusnames, localspluskinds);
     }
+
+    // gh-110543: Make sure the CO_FAST_HIDDEN flag is set correctly.
+    if (!(flags & CO_OPTIMIZED)) {
+        Py_ssize_t code_len = PyBytes_GET_SIZE(code);
+        _Py_CODEUNIT *code_data = (_Py_CODEUNIT *)PyBytes_AS_STRING(code);
+        Py_ssize_t num_code_units = code_len / sizeof(_Py_CODEUNIT);
+        int extended_arg = 0;
+        for (int i = 0; i < num_code_units; i += 1 + _PyOpcode_Caches[code_data[i].op.code]) {
+            _Py_CODEUNIT *instr = &code_data[i];
+            uint8_t opcode = instr->op.code;
+            if (opcode == EXTENDED_ARG) {
+                extended_arg = extended_arg << 8 | instr->op.arg;
+                continue;
+            }
+            if (opcode == LOAD_FAST_AND_CLEAR) {
+                int oparg = extended_arg << 8 | instr->op.arg;
+                if (oparg >= nlocalsplus) {
+                    PyErr_Format(PyExc_ValueError,
+                                "code: LOAD_FAST_AND_CLEAR oparg %d out of range",
+                                oparg);
+                    goto error;
+                }
+                _PyLocals_Kind kind = _PyLocals_GetKind(localspluskinds, oparg);
+                _PyLocals_SetKind(localspluskinds, oparg, kind | CO_FAST_HIDDEN);
+            }
+            extended_arg = 0;
+        }
+    }
+
     // If any cells were args then nlocalsplus will have shrunk.
     if (nlocalsplus != PyTuple_GET_SIZE(localsplusnames)) {
         if (_PyTuple_Resize(&localsplusnames, nlocalsplus) < 0