From e3c83ddddc7cbb890f3e9b367f10af944231001d Mon Sep 17 00:00:00 2001
From: Irit Katriel <iritkatriel@yahoo.com>
Date: Mon, 8 May 2023 20:04:34 +0100
Subject: [PATCH 1/2] gh-104240: return code unit metadata from codegen

---
 Lib/test/support/bytecode_helper.py |  2 +-
 Lib/test/test_compiler_assemble.py  |  2 +-
 Modules/_testinternalcapi.c         |  2 +-
 Python/compile.c                    | 50 +++++++++++++++++++++++++++--
 4 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/Lib/test/support/bytecode_helper.py b/Lib/test/support/bytecode_helper.py
index 357ec44dbc218d..7b577f54b8adc4 100644
--- a/Lib/test/support/bytecode_helper.py
+++ b/Lib/test/support/bytecode_helper.py
@@ -124,7 +124,7 @@ def complete_insts_info(self, insts):
 class CodegenTestCase(CompilationStepTestCase):
 
     def generate_code(self, ast):
-        insts = compiler_codegen(ast, "my_file.py", 0)
+        insts, _ = compiler_codegen(ast, "my_file.py", 0)
         return insts
 
 
diff --git a/Lib/test/test_compiler_assemble.py b/Lib/test/test_compiler_assemble.py
index 96c1691e24a9c8..0bd7a09b001c1d 100644
--- a/Lib/test/test_compiler_assemble.py
+++ b/Lib/test/test_compiler_assemble.py
@@ -52,7 +52,7 @@ def test_simple_expr(self):
             'filename' : 'avg.py',
             'name'     : 'avg',
             'qualname' : 'stats.avg',
-            'consts'   : [2],
+            'consts'   : {2 : 0},
             'argcount' : 2,
             'varnames' : {'x' : 0, 'y' : 1},
         }
diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
index 40ad6f88868daf..8009dca3d0b746 100644
--- a/Modules/_testinternalcapi.c
+++ b/Modules/_testinternalcapi.c
@@ -670,7 +670,7 @@ _testinternalcapi_assemble_code_object_impl(PyObject *module,
     umd.u_cellvars = PyDict_GetItemString(metadata, "cellvars");
     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_varnames));
     assert(PyDict_Check(umd.u_cellvars));
diff --git a/Python/compile.c b/Python/compile.c
index f875e4e17e0a9f..7f87a0f084dd18 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -7261,6 +7261,7 @@ _PyCompile_CodeGen(PyObject *ast, PyObject *filename, PyCompilerFlags *pflags,
                    int optimize, int compile_mode)
 {
     PyObject *res = NULL;
+    PyObject *metadata = NULL;
 
     if (!PyAST_Check(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) {
         goto finally;
     }
+
+    _PyCompile_CodeUnitMetadata *umd = &c->u->u_metadata;
+    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; \
+    }
+
+    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(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) {
-        return NULL;
+        goto finally;
     }
 
-    res = instr_sequence_to_instructions(INSTR_SEQUENCE(c));
+    PyObject *insts = instr_sequence_to_instructions(INSTR_SEQUENCE(c));
+    if (insts == NULL) {
+        goto finally;
+    }
+    res = PyTuple_Pack(2, insts, metadata);
+    Py_DECREF(insts);
 
 finally:
+    Py_XDECREF(metadata);
     compiler_exit_scope(c);
     compiler_free(c);
     _PyArena_Free(arena);
@@ -7375,10 +7415,14 @@ _PyCompile_Assemble(_PyCompile_CodeUnitMetadata *umd, PyObject *filename,
         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,
                                     consts, maxdepth, &optimized_instrs,
                                     nlocalsplus, code_flags, filename);
+    Py_DECREF(consts);
 
 error:
     Py_DECREF(const_cache);

From fc60712b15838b547c1507eda346c647c75f6fc8 Mon Sep 17 00:00:00 2001
From: Irit Katriel <iritkatriel@yahoo.com>
Date: Tue, 9 May 2023 10:39:26 +0100
Subject: [PATCH 2/2] cast to long

---
 Python/compile.c | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Python/compile.c b/Python/compile.c
index 7f87a0f084dd18..12ae831361963d 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -7309,7 +7309,7 @@ _PyCompile_CodeGen(PyObject *ast, PyObject *filename, PyCompilerFlags *pflags,
 #undef SET_MATADATA_ITEM
 
 #define SET_MATADATA_INT(key, value) do { \
-        PyObject *v = PyLong_FromLong(value); \
+        PyObject *v = PyLong_FromLong((long)value); \
         if (v == NULL) goto finally; \
         int res = PyDict_SetItemString(metadata, key, v); \
         Py_XDECREF(v); \