gh-87092: expose the compiler's codegen to python for unit tests

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);

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)

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

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)

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):

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 */