pythongh-104240: make _PyCompile_CodeGen support different compilatio…

…n modes (python#104241)
jbower-fb · May 8, 2023 · c195a2c · c195a2c
1 parent 86ca8d2
commit c195a2c
Show file tree

Hide file tree

Showing 9 changed files with 43 additions and 15 deletions.
diff --git 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
diff --git 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
diff --git a/Include/internal/pycore_unicodeobject_generated.h b/Include/internal/pycore_unicodeobject_generated.h
diff --git a/Lib/test/test_compiler_codegen.py b/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)
diff --git 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
diff --git a/Python/compile.c b/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));