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

diff --git a/Include/internal/pycore_compile.h b/Include/internal/pycore_compile.h
index 4bd4ef57238f9891589e3117dcc73478b2b98389..d2b12c91fe7a00be2e9cd5fbd193f4b3f65776fc 100644 (file)
--- a/Include/internal/pycore_compile.h
+++ b/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,

diff --git a/Include/internal/pycore_global_objects_fini_generated.h b/Include/internal/pycore_global_objects_fini_generated.h
index 9377fd8526e3a2d2f2dde438c2c208b3b9f20901..7e495817981f063bda76534edd21e7579a04b1cb 100644 (file)
--- a/Include/internal/pycore_global_objects_fini_generated.h
+++ b/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));

diff --git a/Include/internal/pycore_global_strings.h b/Include/internal/pycore_global_strings.h
index ed9b2bb44ddffc4adbe4a9108298c389f147bc81..8ebfee85c87c23b4fdb8b4ac8af66dae2afe077a 100644 (file)
--- a/Include/internal/pycore_global_strings.h
+++ b/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)

diff --git a/Include/internal/pycore_runtime_init_generated.h b/Include/internal/pycore_runtime_init_generated.h
index 6ade8fb6eade0381e14df4af3d12fa711104ca20..7b9c73dd1edf3bb26947db0d559425d95f1da1e4 100644 (file)
--- a/Include/internal/pycore_runtime_init_generated.h
+++ b/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), \

diff --git a/Include/internal/pycore_unicodeobject_generated.h b/Include/internal/pycore_unicodeobject_generated.h
index 0b33ea187e60ff68cda62caf76d6f77576d6010d..8e086edbdf81932bebf62d033de57f1141863413 100644 (file)
--- a/Include/internal/pycore_unicodeobject_generated.h
+++ b/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);

diff --git a/Lib/test/test_compiler_codegen.py b/Lib/test/test_compiler_codegen.py
index 022753e0c99483e46db42bb027254500f66881a9..ea57df9cd2400b8bc37166f49edfec6d6ca951f9 100644 (file)
--- a/Lib/test/test_compiler_codegen.py
+++ b/Lib/test/test_compiler_codegen.py
@@ -25,6 +25,8 @@ class IsolatedCodeGenTests(CodegenTestCase):
             ('LOAD_CONST', 2, 1),
             exit_lbl,
             ('POP_TOP', None),
+            ('LOAD_CONST', 3),
+            ('RETURN_VALUE', None),
         ]
         self.codegen_test(snippet, expected)
 
@@ -46,5 +48,7 @@ class IsolatedCodeGenTests(CodegenTestCase):
             ('JUMP', loop_lbl),
             exit_lbl,
             ('END_FOR', None),
+            ('LOAD_CONST', 0),
+            ('RETURN_VALUE', None),
         ]
         self.codegen_test(snippet, expected)

diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
index f35e3b48df9321dd5f70785fd921db140c5db860..40ad6f88868daf7d8cf2403b2314f02b1740cd25 100644 (file)
--- a/Modules/_testinternalcapi.c
+++ b/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);
 }
 
 

diff --git a/Modules/clinic/_testinternalcapi.c.h b/Modules/clinic/_testinternalcapi.c.h
index 8957322257259498f5c431e6e9d8a6d1022d3fd0..41dd50437956c4cea997751c0186ed03ef327411 100644 (file)
--- a/Modules/clinic/_testinternalcapi.c.h
+++ b/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]*/

diff --git a/Python/compile.c b/Python/compile.c
index cbe5403aafbc87bc3b6a64240357b555abdc6aea..f875e4e17e0a9f5aeb04dd7322fb49c68e996ec8 100644 (file)
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -7258,7 +7258,7 @@ error:
 
 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));