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

1 files changed, 6 insertions, 2 deletions

diff --git a/Python/compile.c b/Python/compile.c
index cbe5403..f875e4e 100644
--- 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));