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

2 files changed, 90 insertions, 13 deletions

diff --git a/Lib/test/support/bytecode_helper.py b/Lib/test/support/bytecode_helper.py
index 1d9b889c92..357ec44dbc 100644
--- a/Lib/test/support/bytecode_helper.py
+++ b/Lib/test/support/bytecode_helper.py
@@ -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 @@ class CompilationStepTestCase(unittest.TestCase):
             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 @@ class CodegenTestCase(CompilationStepTestCase):
 
 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
diff --git a/Lib/test/test_compiler_assemble.py b/Lib/test/test_compiler_assemble.py
new file mode 100644
index 0000000000..96c1691e24
--- /dev/null
+++ b/Lib/test/test_compiler_assemble.py
@@ -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)