gh-104240: make _PyCompile_CodeGen support different compilation modes (#104241)

Include/internal/pycore_compile.h

@@ -97,7 +97,8 @@ PyAPI_FUNC(PyObject*) _PyCompile_CodeGen(
         PyObject *ast,
         PyObject *filename,
         PyCompilerFlags *flags,
-        int optimize);
+        int optimize,
+        int compile_mode);
 
 PyAPI_FUNC(PyObject*) _PyCompile_OptimizeCfg(
         PyObject *instructions,

Include/internal/pycore_global_objects_fini_generated.h

@@ -847,6 +847,7 @@ _PyStaticObjects_CheckRefcnt(PyInterpreterState *interp) {
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(code));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(command));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(comment_factory));
+    _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(compile_mode));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(consts));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(context));
     _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(cookie));

Include/internal/pycore_global_strings.h

@@ -335,6 +335,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(code)
         STRUCT_FOR_ID(command)
         STRUCT_FOR_ID(comment_factory)
+        STRUCT_FOR_ID(compile_mode)
         STRUCT_FOR_ID(consts)
         STRUCT_FOR_ID(context)
         STRUCT_FOR_ID(cookie)

Include/internal/pycore_runtime_init_generated.h

@@ -841,6 +841,7 @@ extern "C" {
     INIT_ID(code), \
     INIT_ID(command), \
     INIT_ID(comment_factory), \
+    INIT_ID(compile_mode), \
     INIT_ID(consts), \
     INIT_ID(context), \
     INIT_ID(cookie), \

Include/internal/pycore_unicodeobject_generated.h

@@ -858,6 +858,9 @@ _PyUnicode_InitStaticStrings(PyInterpreterState *interp) {
     string = &_Py_ID(comment_factory);
     assert(_PyUnicode_CheckConsistency(string, 1));
     _PyUnicode_InternInPlace(interp, &string);
+    string = &_Py_ID(compile_mode);
+    assert(_PyUnicode_CheckConsistency(string, 1));
+    _PyUnicode_InternInPlace(interp, &string);
     string = &_Py_ID(consts);
     assert(_PyUnicode_CheckConsistency(string, 1));
     _PyUnicode_InternInPlace(interp, &string);

Lib/test/test_compiler_codegen.py

@@ -25,6 +25,8 @@ def test_if_expression(self):
             ('LOAD_CONST', 2, 1),
             exit_lbl,
             ('POP_TOP', None),
+            ('LOAD_CONST', 3),
+            ('RETURN_VALUE', None),
         ]
         self.codegen_test(snippet, expected)
 
@@ -46,5 +48,7 @@ def test_for_loop(self):
             ('JUMP', loop_lbl),
             exit_lbl,
             ('END_FOR', None),
+            ('LOAD_CONST', 0),
+            ('RETURN_VALUE', None),
         ]
         self.codegen_test(snippet, expected)

Modules/_testinternalcapi.c

@@ -593,17 +593,19 @@ _testinternalcapi.compiler_codegen -> object
   ast: object
   filename: object
   optimize: int
+  compile_mode: int = 0
 
 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]*/
+                                        PyObject *filename, int optimize,
+                                        int compile_mode)
+/*[clinic end generated code: output=40a68f6e13951cc8 input=a0e00784f1517cd7]*/
 {
     PyCompilerFlags *flags = NULL;
-    return _PyCompile_CodeGen(ast, filename, flags, optimize);
+    return _PyCompile_CodeGen(ast, filename, flags, optimize, compile_mode);
 }
 
 

Modules/clinic/_testinternalcapi.c.h

@@ -9,7 +9,7 @@ preserve
 
 
 PyDoc_STRVAR(_testinternalcapi_compiler_codegen__doc__,
-"compiler_codegen($module, /, ast, filename, optimize)\n"
+"compiler_codegen($module, /, ast, filename, optimize, compile_mode=0)\n"
 "--\n"
 "\n"
 "Apply compiler code generation to an AST.");
@@ -19,7 +19,8 @@ PyDoc_STRVAR(_testinternalcapi_compiler_codegen__doc__,
 
 static PyObject *
 _testinternalcapi_compiler_codegen_impl(PyObject *module, PyObject *ast,
-                                        PyObject *filename, int optimize);
+                                        PyObject *filename, int optimize,
+                                        int compile_mode);
 
 static PyObject *
 _testinternalcapi_compiler_codegen(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
@@ -27,14 +28,14 @@ _testinternalcapi_compiler_codegen(PyObject *module, PyObject *const *args, Py_s
     PyObject *return_value = NULL;
     #if defined(Py_BUILD_CORE) && !defined(Py_BUILD_CORE_MODULE)
 
-    #define NUM_KEYWORDS 3
+    #define NUM_KEYWORDS 4
     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), },
+        .ob_item = { &_Py_ID(ast), &_Py_ID(filename), &_Py_ID(optimize), &_Py_ID(compile_mode), },
     };
     #undef NUM_KEYWORDS
     #define KWTUPLE (&_kwtuple.ob_base.ob_base)
@@ -43,19 +44,21 @@ _testinternalcapi_compiler_codegen(PyObject *module, PyObject *const *args, Py_s
     #  define KWTUPLE NULL
     #endif  // !Py_BUILD_CORE
 
-    static const char * const _keywords[] = {"ast", "filename", "optimize", NULL};
+    static const char * const _keywords[] = {"ast", "filename", "optimize", "compile_mode", NULL};
     static _PyArg_Parser _parser = {
         .keywords = _keywords,
         .fname = "compiler_codegen",
         .kwtuple = KWTUPLE,
     };
     #undef KWTUPLE
-    PyObject *argsbuf[3];
+    PyObject *argsbuf[4];
+    Py_ssize_t noptargs = nargs + (kwnames ? PyTuple_GET_SIZE(kwnames) : 0) - 3;
     PyObject *ast;
     PyObject *filename;
     int optimize;
+    int compile_mode = 0;
 
-    args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 3, 3, 0, argsbuf);
+    args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 3, 4, 0, argsbuf);
     if (!args) {
         goto exit;
     }
@@ -65,7 +68,15 @@ _testinternalcapi_compiler_codegen(PyObject *module, PyObject *const *args, Py_s
     if (optimize == -1 && PyErr_Occurred()) {
         goto exit;
     }
-    return_value = _testinternalcapi_compiler_codegen_impl(module, ast, filename, optimize);
+    if (!noptargs) {
+        goto skip_optional_pos;
+    }
+    compile_mode = _PyLong_AsInt(args[3]);
+    if (compile_mode == -1 && PyErr_Occurred()) {
+        goto exit;
+    }
+skip_optional_pos:
+    return_value = _testinternalcapi_compiler_codegen_impl(module, ast, filename, optimize, compile_mode);
 
 exit:
     return return_value;
@@ -190,4 +201,4 @@ _testinternalcapi_assemble_code_object(PyObject *module, PyObject *const *args,
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=d5e08c9d67f9721f input=a9049054013a1b77]*/
+/*[clinic end generated code: output=ab661d56a14b1a1c input=a9049054013a1b77]*/

Python/compile.c

@@ -7258,7 +7258,7 @@ cfg_to_instructions(cfg_builder *g)
 
 PyObject *
 _PyCompile_CodeGen(PyObject *ast, PyObject *filename, PyCompilerFlags *pflags,
-                   int optimize)
+                   int optimize, int compile_mode)
 {
     PyObject *res = NULL;
 
@@ -7272,7 +7272,7 @@ _PyCompile_CodeGen(PyObject *ast, PyObject *filename, PyCompilerFlags *pflags,
         return NULL;
     }
 
-    mod_ty mod = PyAST_obj2mod(ast, arena, 0 /* exec */);
+    mod_ty mod = PyAST_obj2mod(ast, arena, compile_mode);
     if (mod == NULL || !_PyAST_Validate(mod)) {
         _PyArena_Free(arena);
         return NULL;
@@ -7287,6 +7287,10 @@ _PyCompile_CodeGen(PyObject *ast, PyObject *filename, PyCompilerFlags *pflags,
     if (compiler_codegen(c, mod) < 0) {
         goto finally;
     }
+    int addNone = mod->kind != Expression_kind;
+    if (add_return_at_end(c, addNone) < 0) {
+        return NULL;
+    }
 
     res = instr_sequence_to_instructions(INSTR_SEQUENCE(c));