python · iritkatriel · Nov 14, 2022 · Nov 4, 2022 · Nov 4, 2022 · Nov 5, 2022
diff --git a/Include/internal/pycore_compile.h b/Include/internal/pycore_compile.h
@@ -40,6 +40,13 @@ extern int _PyAST_Optimize(
     _PyASTOptimizeState *state);
 
 /* Access compiler internals for unit testing */
+
+PyAPI_FUNC(PyObject*) _PyCompile_CodeGen(
+        PyObject *ast,
+        PyObject *filename,
+        PyCompilerFlags *flags,
+        int optimize);
+
 PyAPI_FUNC(PyObject*) _PyCompile_OptimizeCfg(
         PyObject *instructions,
         PyObject *consts);

diff --git a/Include/internal/pycore_global_objects_fini_generated.h b/Include/internal/pycore_global_objects_fini_generated.h
diff --git a/Include/internal/pycore_global_strings.h b/Include/internal/pycore_global_strings.h
@@ -267,6 +267,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(arguments)
         STRUCT_FOR_ID(argv)
         STRUCT_FOR_ID(as_integer_ratio)
+        STRUCT_FOR_ID(ast)
         STRUCT_FOR_ID(attribute)
         STRUCT_FOR_ID(authorizer_callback)
         STRUCT_FOR_ID(autocommit)

diff --git a/Include/internal/pycore_runtime_init_generated.h b/Include/internal/pycore_runtime_init_generated.h
diff --git a/Include/internal/pycore_unicodeobject_generated.h b/Include/internal/pycore_unicodeobject_generated.h
diff --git 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 optimize_cfg
+from _testinternalcapi import compiler_codegen, optimize_cfg
 
 _UNSPECIFIED = object()
 
@@ -44,8 +44,7 @@ def assertNotInBytecode(self, x, opname, argval=_UNSPECIFIED):
                     msg = msg % (opname, argval, disassembly)
                     self.fail(msg)
 
-
-class CfgOptimizationTestCase(unittest.TestCase):
+class CompilationStepTestCase(unittest.TestCase):
 
     HAS_ARG = set(dis.hasarg)
     HAS_TARGET = set(dis.hasjrel + dis.hasjabs + dis.hasexc)
@@ -58,24 +57,35 @@ def Label(self):
         self.last_label += 1
         return self.last_label
 
-    def complete_insts_info(self, insts):
-        # fill in omitted fields in location, and oparg 0 for ops with no arg.
-        instructions = []
-        for item in insts:
-            if isinstance(item, int):
-                instructions.append(item)
-            else:
-                assert isinstance(item, tuple)
-                inst = list(reversed(item))
-                opcode = dis.opmap[inst.pop()]
-                oparg = inst.pop() if opcode in self.HAS_ARG_OR_TARGET else 0
-                loc = inst + [-1] * (4 - len(inst))
-                instructions.append((opcode, oparg, *loc))
-        return instructions
+    def assertInstructionsMatch(self, actual_, expected_):
+        # get two lists where each entry is a label or
+        # an instruction tuple. Compare them, while mapping
+        # each actual label to a corresponding expected label
+        # based on their locations.
+
+        self.assertIsInstance(actual_, list)
+        self.assertIsInstance(expected_, list)
+
+        actual = self.normalize_insts(actual_)
+        expected = self.normalize_insts(expected_)
+        self.assertEqual(len(actual), len(expected))
+
+        # compare instructions
+        for act, exp in zip(actual, expected):
+            if isinstance(act, int):
+                self.assertEqual(exp, act)
+                continue
+            self.assertIsInstance(exp, tuple)
+            self.assertIsInstance(act, tuple)
+            # crop comparison to the provided expected values
+            if len(act) > len(exp):
+                act = act[:len(exp)]
+            self.assertEqual(exp, act)
 
     def normalize_insts(self, insts):
         """ Map labels to instruction index.
             Remove labels which are not used as jump targets.
+            Map opcodes to opnames.
         """
         labels_map = {}
         targets = set()
@@ -107,31 +117,32 @@ def normalize_insts(self, insts):
                 res.append((opcode, arg, *loc))
         return res
 
-    def get_optimized(self, insts, consts):
-        insts = self.complete_insts_info(insts)
-        insts = optimize_cfg(insts, consts)
-        return insts, consts
 
-    def compareInstructions(self, actual_, expected_):
-        # get two lists where each entry is a label or
-        # an instruction tuple. Compare them, while mapping
-        # each actual label to a corresponding expected label
-        # based on their locations.
+class CodegenTestCase(CompilationStepTestCase):
 
-        self.assertIsInstance(actual_, list)
-        self.assertIsInstance(expected_, list)
+    def generate_code(self, ast):
+        insts = compiler_codegen(ast, "my_file.py", 0)
+        return insts
 
-        actual = self.normalize_insts(actual_)
-        expected = self.normalize_insts(expected_)
-        self.assertEqual(len(actual), len(expected))
 
-        # compare instructions
-        for act, exp in zip(actual, expected):
-            if isinstance(act, int):
-                self.assertEqual(exp, act)
-                continue
-            self.assertIsInstance(exp, tuple)
-            self.assertIsInstance(act, tuple)
-            # pad exp with -1's (if location info is incomplete)
-            exp += (-1,) * (len(act) - len(exp))
-            self.assertEqual(exp, act)
+class CfgOptimizationTestCase(CompilationStepTestCase):
+
+    def complete_insts_info(self, insts):
+        # fill in omitted fields in location, and oparg 0 for ops with no arg.
+        instructions = []
+        for item in insts:
+            if isinstance(item, int):
+                instructions.append(item)
+            else:
+                assert isinstance(item, tuple)
+                inst = list(reversed(item))
+                opcode = dis.opmap[inst.pop()]
+                oparg = inst.pop() if opcode in self.HAS_ARG_OR_TARGET else 0
+                loc = inst + [-1] * (4 - len(inst))
+                instructions.append((opcode, oparg, *loc))
+        return instructions
+
+    def get_optimized(self, insts, consts):
+        insts = self.complete_insts_info(insts)
+        insts = optimize_cfg(insts, consts)
+        return insts, consts
diff --git a/Lib/test/test_compiler_codegen.py b/Lib/test/test_compiler_codegen.py
@@ -0,0 +1,50 @@
+
+from test.support.bytecode_helper import CodegenTestCase
+
+# Tests for the code-generation stage of the compiler.
+# Examine the un-optimized code generated from the AST.
+
+class IsolatedCodeGenTests(CodegenTestCase):
+
+    def codegen_test(self, snippet, expected_insts):
+        import ast
+        a = ast.parse(snippet, "my_file.py", "exec");
+        insts = self.generate_code(a)
+        self.assertInstructionsMatch(insts, expected_insts)
+
+    def test_if_expression(self):
+        snippet = "42 if True else 24"
+        false_lbl = self.Label()
+        expected = [
+            ('RESUME', 0, 0),
+            ('LOAD_CONST', 0, 1),
+            ('POP_JUMP_IF_FALSE', false_lbl := self.Label(), 1),
+            ('LOAD_CONST', 1, 1),
+            ('JUMP', exit_lbl := self.Label()),
+            false_lbl,
+            ('LOAD_CONST', 2, 1),
+            exit_lbl,
+            ('POP_TOP', None),
+        ]
+        self.codegen_test(snippet, expected)
+
+    def test_for_loop(self):
+        snippet = "for x in l:\n\tprint(x)"
+        false_lbl = self.Label()
+        expected = [
+            ('RESUME', 0, 0),
+            ('LOAD_NAME', 0, 1),
+            ('GET_ITER', None, 1),
+            loop_lbl := self.Label(),
+            ('FOR_ITER', exit_lbl := self.Label(), 1),
+            ('STORE_NAME', None, 1),
+            ('PUSH_NULL', None, 2),
+            ('LOAD_NAME', None, 2),
+            ('LOAD_NAME', None, 2),
+            ('CALL', None, 2),
+            ('POP_TOP', None),
+            ('JUMP', loop_lbl),
+            exit_lbl,
+            ('END_FOR', None),
+        ]
+        self.codegen_test(snippet, expected)
diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py
@@ -984,7 +984,7 @@ def cfg_optimization_test(self, insts, expected_insts,
         if expected_consts is None:
             expected_consts = consts
         opt_insts, opt_consts = self.get_optimized(insts, consts)
-        self.compareInstructions(opt_insts, expected_insts)
+        self.assertInstructionsMatch(opt_insts, expected_insts)
         self.assertEqual(opt_consts, expected_consts)
 
     def test_conditional_jump_forward_non_const_condition(self):

diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
@@ -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_OptimizeCfg()
+#include "pycore_compile.h"      // _PyCompile_CodeGen, _PyCompile_OptimizeCfg
 #include "pycore_fileutils.h"    // _Py_normpath
 #include "pycore_frame.h"        // _PyInterpreterFrame
 #include "pycore_gc.h"           // PyGC_Head
@@ -529,6 +529,26 @@ set_eval_frame_record(PyObject *self, PyObject *list)
     Py_RETURN_NONE;
 }
 
+/*[clinic input]
+
+_testinternalcapi.compiler_codegen -> object
+
+  ast: object
+  filename: object
+  optimize: int
+
+Apply compiler code generation to an AST.
+[clinic start generated code]*/
+
+static PyObject *
+_testinternalcapi_compiler_codegen_impl(PyObject *module, PyObject *ast,
+                                        PyObject *filename, int optimize)
+/*[clinic end generated code: output=fbbbbfb34700c804 input=e9fbe6562f7f75e4]*/
+{
+    PyCompilerFlags *flags = NULL;
+    return _PyCompile_CodeGen(ast, filename, flags, optimize);
+}
+
 
 /*[clinic input]
 
@@ -612,6 +632,7 @@ static PyMethodDef TestMethods[] = {
     {"DecodeLocaleEx", decode_locale_ex, METH_VARARGS},
     {"set_eval_frame_default", set_eval_frame_default, METH_NOARGS, NULL},
     {"set_eval_frame_record", set_eval_frame_record, METH_O, NULL},
+    _TESTINTERNALCAPI_COMPILER_CODEGEN_METHODDEF
     _TESTINTERNALCAPI_OPTIMIZE_CFG_METHODDEF
     {"get_interp_settings", get_interp_settings, METH_VARARGS, NULL},
     {NULL, NULL} /* sentinel */

diff --git a/Modules/clinic/_testinternalcapi.c.h b/Modules/clinic/_testinternalcapi.c.h