gh-104240: return code unit metadata from codegen (#104300)

This commit is contained in:
Irit Katriel 2023-05-09 14:33:40 +01:00 committed by GitHub
parent c21f828760
commit ca95edf177
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 52 additions and 8 deletions

View file

@ -124,7 +124,7 @@ class CompilationStepTestCase(unittest.TestCase):
class CodegenTestCase(CompilationStepTestCase): class CodegenTestCase(CompilationStepTestCase):
def generate_code(self, ast): def generate_code(self, ast):
insts = compiler_codegen(ast, "my_file.py", 0) insts, _ = compiler_codegen(ast, "my_file.py", 0)
return insts return insts

View file

@ -52,7 +52,7 @@ class IsolatedAssembleTests(AssemblerTestCase):
'filename' : 'avg.py', 'filename' : 'avg.py',
'name' : 'avg', 'name' : 'avg',
'qualname' : 'stats.avg', 'qualname' : 'stats.avg',
'consts' : [2], 'consts' : {2 : 0},
'argcount' : 2, 'argcount' : 2,
'varnames' : {'x' : 0, 'y' : 1}, 'varnames' : {'x' : 0, 'y' : 1},
} }

View file

@ -670,7 +670,7 @@ _testinternalcapi_assemble_code_object_impl(PyObject *module,
umd.u_cellvars = PyDict_GetItemString(metadata, "cellvars"); umd.u_cellvars = PyDict_GetItemString(metadata, "cellvars");
umd.u_freevars = PyDict_GetItemString(metadata, "freevars"); umd.u_freevars = PyDict_GetItemString(metadata, "freevars");
assert(PyList_Check(umd.u_consts)); assert(PyDict_Check(umd.u_consts));
assert(PyDict_Check(umd.u_names)); assert(PyDict_Check(umd.u_names));
assert(PyDict_Check(umd.u_varnames)); assert(PyDict_Check(umd.u_varnames));
assert(PyDict_Check(umd.u_cellvars)); assert(PyDict_Check(umd.u_cellvars));

View file

@ -7261,6 +7261,7 @@ _PyCompile_CodeGen(PyObject *ast, PyObject *filename, PyCompilerFlags *pflags,
int optimize, int compile_mode) int optimize, int compile_mode)
{ {
PyObject *res = NULL; PyObject *res = NULL;
PyObject *metadata = NULL;
if (!PyAST_Check(ast)) { if (!PyAST_Check(ast)) {
PyErr_SetString(PyExc_TypeError, "expected an AST"); PyErr_SetString(PyExc_TypeError, "expected an AST");
@ -7287,14 +7288,53 @@ _PyCompile_CodeGen(PyObject *ast, PyObject *filename, PyCompilerFlags *pflags,
if (compiler_codegen(c, mod) < 0) { if (compiler_codegen(c, mod) < 0) {
goto finally; goto finally;
} }
int addNone = mod->kind != Expression_kind;
if (add_return_at_end(c, addNone) < 0) { _PyCompile_CodeUnitMetadata *umd = &c->u->u_metadata;
return NULL; metadata = PyDict_New();
if (metadata == NULL) {
goto finally;
}
#define SET_MATADATA_ITEM(key, value) \
if (value != NULL) { \
if (PyDict_SetItemString(metadata, key, value) < 0) goto finally; \
} }
res = instr_sequence_to_instructions(INSTR_SEQUENCE(c)); SET_MATADATA_ITEM("name", umd->u_name);
SET_MATADATA_ITEM("qualname", umd->u_qualname);
SET_MATADATA_ITEM("consts", umd->u_consts);
SET_MATADATA_ITEM("names", umd->u_names);
SET_MATADATA_ITEM("varnames", umd->u_varnames);
SET_MATADATA_ITEM("cellvars", umd->u_cellvars);
SET_MATADATA_ITEM("freevars", umd->u_freevars);
#undef SET_MATADATA_ITEM
#define SET_MATADATA_INT(key, value) do { \
PyObject *v = PyLong_FromLong((long)value); \
if (v == NULL) goto finally; \
int res = PyDict_SetItemString(metadata, key, v); \
Py_XDECREF(v); \
if (res < 0) goto finally; \
} while (0);
SET_MATADATA_INT("argcount", umd->u_argcount);
SET_MATADATA_INT("posonlyargcount", umd->u_posonlyargcount);
SET_MATADATA_INT("kwonlyargcount", umd->u_kwonlyargcount);
#undef SET_MATADATA_INT
int addNone = mod->kind != Expression_kind;
if (add_return_at_end(c, addNone) < 0) {
goto finally;
}
PyObject *insts = instr_sequence_to_instructions(INSTR_SEQUENCE(c));
if (insts == NULL) {
goto finally;
}
res = PyTuple_Pack(2, insts, metadata);
Py_DECREF(insts);
finally: finally:
Py_XDECREF(metadata);
compiler_exit_scope(c); compiler_exit_scope(c);
compiler_free(c); compiler_free(c);
_PyArena_Free(arena); _PyArena_Free(arena);
@ -7375,10 +7415,14 @@ _PyCompile_Assemble(_PyCompile_CodeUnitMetadata *umd, PyObject *filename,
goto error; goto error;
} }
PyObject *consts = umd->u_consts; PyObject *consts = consts_dict_keys_inorder(umd->u_consts);
if (consts == NULL) {
goto error;
}
co = _PyAssemble_MakeCodeObject(umd, const_cache, co = _PyAssemble_MakeCodeObject(umd, const_cache,
consts, maxdepth, &optimized_instrs, consts, maxdepth, &optimized_instrs,
nlocalsplus, code_flags, filename); nlocalsplus, code_flags, filename);
Py_DECREF(consts);
error: error:
Py_DECREF(const_cache); Py_DECREF(const_cache);