gh-87092: expose the compiler's codegen to python for unit tests (GH-99111)

2 files changed, 86 insertions, 2 deletions

diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
index 83ce756..cec114c 100644
--- 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
index 8113fff..e8d5681 100644
--- a/Modules/clinic/_testinternalcapi.c.h
+++ b/Modules/clinic/_testinternalcapi.c.h
@@ -8,6 +8,69 @@ preserve
 #endif
 
 
+PyDoc_STRVAR(_testinternalcapi_compiler_codegen__doc__,
+"compiler_codegen($module, /, ast, filename, optimize)\n"
+"--\n"
+"\n"
+"Apply compiler code generation to an AST.");
+
+#define _TESTINTERNALCAPI_COMPILER_CODEGEN_METHODDEF    \
+    {"compiler_codegen", _PyCFunction_CAST(_testinternalcapi_compiler_codegen), METH_FASTCALL|METH_KEYWORDS, _testinternalcapi_compiler_codegen__doc__},
+
+static PyObject *
+_testinternalcapi_compiler_codegen_impl(PyObject *module, PyObject *ast,
+                                        PyObject *filename, int optimize);
+
+static PyObject *
+_testinternalcapi_compiler_codegen(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
+{
+    PyObject *return_value = NULL;
+    #if defined(Py_BUILD_CORE) && !defined(Py_BUILD_CORE_MODULE)
+
+    #define NUM_KEYWORDS 3
+    static struct {
+        PyGC_Head _this_is_not_used;
+        PyObject_VAR_HEAD
+        PyObject *ob_item[NUM_KEYWORDS];
+    } _kwtuple = {
+        .ob_base = PyVarObject_HEAD_INIT(&PyTuple_Type, NUM_KEYWORDS)
+        .ob_item = { &_Py_ID(ast), &_Py_ID(filename), &_Py_ID(optimize), },
+    };
+    #undef NUM_KEYWORDS
+    #define KWTUPLE (&_kwtuple.ob_base.ob_base)
+
+    #else  // !Py_BUILD_CORE
+    #  define KWTUPLE NULL
+    #endif  // !Py_BUILD_CORE
+
+    static const char * const _keywords[] = {"ast", "filename", "optimize", NULL};
+    static _PyArg_Parser _parser = {
+        .keywords = _keywords,
+        .fname = "compiler_codegen",
+        .kwtuple = KWTUPLE,
+    };
+    #undef KWTUPLE
+    PyObject *argsbuf[3];
+    PyObject *ast;
+    PyObject *filename;
+    int optimize;
+
+    args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 3, 3, 0, argsbuf);
+    if (!args) {
+        goto exit;
+    }
+    ast = args[0];
+    filename = args[1];
+    optimize = _PyLong_AsInt(args[2]);
+    if (optimize == -1 && PyErr_Occurred()) {
+        goto exit;
+    }
+    return_value = _testinternalcapi_compiler_codegen_impl(module, ast, filename, optimize);
+
+exit:
+    return return_value;
+}
+
 PyDoc_STRVAR(_testinternalcapi_optimize_cfg__doc__,
 "optimize_cfg($module, /, instructions, consts)\n"
 "--\n"
@@ -65,4 +128,4 @@ _testinternalcapi_optimize_cfg(PyObject *module, PyObject *const *args, Py_ssize
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=3b1fd713290f68a9 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=efe95836482fd542 input=a9049054013a1b77]*/