diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py index 12f7bbd123b30c..f95a78aff0c711 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 @@ 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 @@ def get_output(moddict, name): 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 @@ def test_frame_locals(self): 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, "", "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 00000000000000..5f9571566da18d --- /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 79ac5749037647..dc46b773c26528 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