gh-87092: Expose assembler to unit tests (#103988)

Author: iritkatriel

Include/internal/pycore_compile.h

@@ -103,6 +103,10 @@ PyAPI_FUNC(PyObject*) _PyCompile_OptimizeCfg(
         PyObject *instructions,
         PyObject *consts);
 
+PyAPI_FUNC(PyCodeObject*)
+_PyCompile_Assemble(_PyCompile_CodeUnitMetadata *umd, PyObject *filename,
+                    PyObject *instructions);
+
 #ifdef __cplusplus
 }
 #endif

Include/internal/pycore_global_objects_fini_generated.h

@@ -517,6 +517,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(memlimit)
         STRUCT_FOR_ID(message)
         STRUCT_FOR_ID(metaclass)
+        STRUCT_FOR_ID(metadata)
         STRUCT_FOR_ID(method)
         STRUCT_FOR_ID(mod)
         STRUCT_FOR_ID(mode)

Include/internal/pycore_global_strings.h

@@ -3,7 +3,7 @@
 import unittest
 import dis
 import io
-from _testinternalcapi import compiler_codegen, optimize_cfg
+from _testinternalcapi import compiler_codegen, optimize_cfg, assemble_code_object
 
 _UNSPECIFIED = object()
 
@@ -108,6 +108,18 @@ def normalize_insts(self, insts):
             res.append((opcode, arg, *loc))
         return res
 
+    def complete_insts_info(self, insts):
+        # fill in omitted fields in location, and oparg 0 for ops with no arg.
+        res = []
+        for item in insts:
+            assert isinstance(item, tuple)
+            inst = list(item)
+            opcode = dis.opmap[inst[0]]
+            oparg = inst[1]
+            loc = inst[2:] + [-1] * (6 - len(inst))
+            res.append((opcode, oparg, *loc))
+        return res
+
 
 class CodegenTestCase(CompilationStepTestCase):
 
@@ -118,20 +130,14 @@ def generate_code(self, ast):
 
 class CfgOptimizationTestCase(CompilationStepTestCase):
 
-    def complete_insts_info(self, insts):
-        # fill in omitted fields in location, and oparg 0 for ops with no arg.
-        res = []
-        for item in insts:
-            assert isinstance(item, tuple)
-            inst = list(reversed(item))
-            opcode = dis.opmap[inst.pop()]
-            oparg = inst.pop()
-            loc = inst + [-1] * (4 - len(inst))
-            res.append((opcode, oparg, *loc))
-        return res
-
     def get_optimized(self, insts, consts):
         insts = self.normalize_insts(insts)
         insts = self.complete_insts_info(insts)
         insts = optimize_cfg(insts, consts)
         return insts, consts
+
+class AssemblerTestCase(CompilationStepTestCase):
+
+    def get_code_object(self, filename, insts, metadata):
+        co = assemble_code_object(filename, insts, metadata)
+        return co

Include/internal/pycore_runtime_init_generated.h

@@ -0,0 +1,71 @@
+
+import ast
+import types
+
+from test.support.bytecode_helper import AssemblerTestCase
+
+
+# Tests for the code-object creation stage of the compiler.
+
+class IsolatedAssembleTests(AssemblerTestCase):
+
+    def complete_metadata(self, metadata, filename="myfile.py"):
+        if metadata is None:
+            metadata = {}
+        for key in ['name', 'qualname']:
+            metadata.setdefault(key, key)
+        for key in ['consts']:
+            metadata.setdefault(key, [])
+        for key in ['names', 'varnames', 'cellvars', 'freevars']:
+            metadata.setdefault(key, {})
+        for key in ['argcount', 'posonlyargcount', 'kwonlyargcount']:
+            metadata.setdefault(key, 0)
+        metadata.setdefault('firstlineno', 1)
+        metadata.setdefault('filename', filename)
+        return metadata
+
+    def assemble_test(self, insts, metadata, expected):
+        metadata = self.complete_metadata(metadata)
+        insts = self.complete_insts_info(insts)
+
+        co = self.get_code_object(metadata['filename'], insts, metadata)
+        self.assertIsInstance(co, types.CodeType)
+
+        expected_metadata = {}
+        for key, value in metadata.items():
+            if isinstance(value, list):
+                expected_metadata[key] = tuple(value)
+            elif isinstance(value, dict):
+                expected_metadata[key] = tuple(value.keys())
+            else:
+                expected_metadata[key] = value
+
+        for key, value in expected_metadata.items():
+            self.assertEqual(getattr(co, "co_" + key), value)
+
+        f = types.FunctionType(co, {})
+        for args, res in expected.items():
+            self.assertEqual(f(*args), res)
+
+    def test_simple_expr(self):
+        metadata = {
+            'filename' : 'avg.py',
+            'name'     : 'avg',
+            'qualname' : 'stats.avg',
+            'consts'   : [2],
+            'argcount' : 2,
+            'varnames' : {'x' : 0, 'y' : 1},
+        }
+
+        # code for "return (x+y)/2"
+        insts = [
+            ('RESUME', 0),
+            ('LOAD_FAST', 0, 1),   # 'x'
+            ('LOAD_FAST', 1, 1),   # 'y'
+            ('BINARY_OP', 0, 1),   # '+'
+            ('LOAD_CONST', 0, 1),  # 2
+            ('BINARY_OP', 11, 1),   # '/'
+            ('RETURN_VALUE', 1),
+        ]
+        expected = {(3, 4) : 3.5, (-100, 200) : 50, (10, 18) : 14}
+        self.assemble_test(insts, metadata, expected)

Include/internal/pycore_unicodeobject_generated.h

@@ -14,7 +14,7 @@
 #include "Python.h"
 #include "pycore_atomic_funcs.h" // _Py_atomic_int_get()
 #include "pycore_bitutils.h"     // _Py_bswap32()
-#include "pycore_compile.h"      // _PyCompile_CodeGen, _PyCompile_OptimizeCfg
+#include "pycore_compile.h"      // _PyCompile_CodeGen, _PyCompile_OptimizeCfg, _PyCompile_Assemble
 #include "pycore_fileutils.h"    // _Py_normpath
 #include "pycore_frame.h"        // _PyInterpreterFrame
 #include "pycore_gc.h"           // PyGC_Head
@@ -625,6 +625,68 @@ _testinternalcapi_optimize_cfg_impl(PyObject *module, PyObject *instructions,
     return _PyCompile_OptimizeCfg(instructions, consts);
 }
 
+static int
+get_nonnegative_int_from_dict(PyObject *dict, const char *key) {
+    PyObject *obj = PyDict_GetItemString(dict, key);
+    if (obj == NULL) {
+        return -1;
+    }
+    return PyLong_AsLong(obj);
+}
+
+/*[clinic input]
+
+_testinternalcapi.assemble_code_object -> object
+
+  filename: object
+  instructions: object
+  metadata: object
+
+Create a code object for the given instructions.
+[clinic start generated code]*/
+
+static PyObject *
+_testinternalcapi_assemble_code_object_impl(PyObject *module,
+                                            PyObject *filename,
+                                            PyObject *instructions,
+                                            PyObject *metadata)
+/*[clinic end generated code: output=38003dc16a930f48 input=e713ad77f08fb3a8]*/
+
+{
+    assert(PyDict_Check(metadata));
+    _PyCompile_CodeUnitMetadata umd;
+
+    umd.u_name = PyDict_GetItemString(metadata, "name");
+    umd.u_qualname = PyDict_GetItemString(metadata, "qualname");
+
+    assert(PyUnicode_Check(umd.u_name));
+    assert(PyUnicode_Check(umd.u_qualname));
+
+    umd.u_consts = PyDict_GetItemString(metadata, "consts");
+    umd.u_names = PyDict_GetItemString(metadata, "names");
+    umd.u_varnames = PyDict_GetItemString(metadata, "varnames");
+    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_names));
+    assert(PyDict_Check(umd.u_varnames));
+    assert(PyDict_Check(umd.u_cellvars));
+    assert(PyDict_Check(umd.u_freevars));
+
+    umd.u_argcount = get_nonnegative_int_from_dict(metadata, "argcount");
+    umd.u_posonlyargcount = get_nonnegative_int_from_dict(metadata, "posonlyargcount");
+    umd.u_kwonlyargcount = get_nonnegative_int_from_dict(metadata, "kwonlyargcount");
+    umd.u_firstlineno = get_nonnegative_int_from_dict(metadata, "firstlineno");
+
+    assert(umd.u_argcount >= 0);
+    assert(umd.u_posonlyargcount >= 0);
+    assert(umd.u_kwonlyargcount >= 0);
+    assert(umd.u_firstlineno >= 0);
+
+    return (PyObject*)_PyCompile_Assemble(&umd, filename, instructions);
+}
+
 
 static PyObject *
 get_interp_settings(PyObject *self, PyObject *args)
@@ -705,6 +767,7 @@ static PyMethodDef module_functions[] = {
     {"set_eval_frame_record", set_eval_frame_record, METH_O, NULL},
     _TESTINTERNALCAPI_COMPILER_CODEGEN_METHODDEF
     _TESTINTERNALCAPI_OPTIMIZE_CFG_METHODDEF
+    _TESTINTERNALCAPI_ASSEMBLE_CODE_OBJECT_METHODDEF
     {"get_interp_settings", get_interp_settings, METH_VARARGS, NULL},
     {"clear_extension", clear_extension, METH_VARARGS, NULL},
     {NULL, NULL} /* sentinel */