gh-108113: Make it possible to optimize an AST (#108282)

Author: iritkatriel

Include/internal/pycore_compile.h

@@ -19,6 +19,14 @@ PyAPI_FUNC(PyCodeObject*) _PyAST_Compile(
     int optimize,
     struct _arena *arena);
 
+/* AST optimizations */
+PyAPI_FUNC(int) _PyCompile_AstOptimize(
+    struct _mod *mod,
+    PyObject *filename,
+    PyCompilerFlags *flags,
+    int optimize,
+    struct _arena *arena);
+
 static const _PyCompilerSrcLocation NO_LOCATION = {-1, -1, -1, -1};
 
 extern int _PyAST_Optimize(

Lib/test/test_ast.py

@@ -361,14 +361,16 @@ def test_optimization_levels__debug__(self):
         cases = [(-1, '__debug__'), (0, '__debug__'), (1, False), (2, False)]
         for (optval, expected) in cases:
             with self.subTest(optval=optval, expected=expected):
-                res = ast.parse("__debug__", optimize=optval)
-                self.assertIsInstance(res.body[0], ast.Expr)
-                if isinstance(expected, bool):
-                    self.assertIsInstance(res.body[0].value, ast.Constant)
-                    self.assertEqual(res.body[0].value.value, expected)
-                else:
-                    self.assertIsInstance(res.body[0].value, ast.Name)
-                    self.assertEqual(res.body[0].value.id, expected)
+                res1 = ast.parse("__debug__", optimize=optval)
+                res2 = ast.parse(ast.parse("__debug__"), optimize=optval)
+                for res in [res1, res2]:
+                    self.assertIsInstance(res.body[0], ast.Expr)
+                    if isinstance(expected, bool):
+                        self.assertIsInstance(res.body[0].value, ast.Constant)
+                        self.assertEqual(res.body[0].value.value, expected)
+                    else:
+                        self.assertIsInstance(res.body[0].value, ast.Name)
+                        self.assertEqual(res.body[0].value.id, expected)
 
     def test_optimization_levels_const_folding(self):
         folded = ('Expr', (1, 0, 1, 5), ('Constant', (1, 0, 1, 5), 3, None))
@@ -381,9 +383,11 @@ def test_optimization_levels_const_folding(self):
         cases = [(-1, not_folded), (0, not_folded), (1, folded), (2, folded)]
         for (optval, expected) in cases:
             with self.subTest(optval=optval):
-                tree = ast.parse("1 + 2", optimize=optval)
-                res = to_tuple(tree.body[0])
-                self.assertEqual(res, expected)
+                tree1 = ast.parse("1 + 2", optimize=optval)
+                tree2 = ast.parse(ast.parse("1 + 2"), optimize=optval)
+                for tree in [tree1, tree2]:
+                    res = to_tuple(tree.body[0])
+                    self.assertEqual(res, expected)
 
     def test_invalid_position_information(self):
         invalid_linenos = [

Lib/test/test_builtin.py

@@ -521,9 +521,10 @@ def test_compile_async_generator(self):
     def test_compile_ast(self):
         args = ("a*(1+2)", "f.py", "exec")
         raw = compile(*args, flags = ast.PyCF_ONLY_AST).body[0]
-        opt = compile(*args, flags = ast.PyCF_OPTIMIZED_AST).body[0]
+        opt1 = compile(*args, flags = ast.PyCF_OPTIMIZED_AST).body[0]
+        opt2 = compile(ast.parse(args[0]), *args[1:], flags = ast.PyCF_OPTIMIZED_AST).body[0]
 
-        for tree in (raw, opt):
+        for tree in (raw, opt1, opt2):
             self.assertIsInstance(tree.value, ast.BinOp)
             self.assertIsInstance(tree.value.op, ast.Mult)
             self.assertIsInstance(tree.value.left, ast.Name)
@@ -536,9 +537,10 @@ def test_compile_ast(self):
         self.assertIsInstance(raw_right.right, ast.Constant)
         self.assertEqual(raw_right.right.value, 2)
 
-        opt_right = opt.value.right  # expect Constant(3)
-        self.assertIsInstance(opt_right, ast.Constant)
-        self.assertEqual(opt_right.value, 3)
+        for opt in [opt1, opt2]:
+            opt_right = opt.value.right  # expect Constant(3)
+            self.assertIsInstance(opt_right, ast.Constant)
+            self.assertEqual(opt_right.value, 3)
 
     def test_delattr(self):
         sys.spam = 1

Python/bltinmodule.c

@@ -804,23 +804,40 @@ builtin_compile_impl(PyObject *module, PyObject *source, PyObject *filename,
     if (is_ast == -1)
         goto error;
     if (is_ast) {
-        if (flags & PyCF_ONLY_AST) {
+        if ((flags & PyCF_OPTIMIZED_AST) == PyCF_ONLY_AST) {
+            // return an un-optimized AST
             result = Py_NewRef(source);
         }
         else {
-            PyArena *arena;
-            mod_ty mod;
+            // Return an optimized AST or code object
 
-            arena = _PyArena_New();
-            if (arena == NULL)
-                goto error;
-            mod = PyAST_obj2mod(source, arena, compile_mode);
-            if (mod == NULL || !_PyAST_Validate(mod)) {
-                _PyArena_Free(arena);
+            PyArena *arena = _PyArena_New();
+            if (arena == NULL) {
                 goto error;
             }
-            result = (PyObject*)_PyAST_Compile(mod, filename,
-                                               &cf, optimize, arena);
+
+            if (flags & PyCF_ONLY_AST) {
+                mod_ty mod = PyAST_obj2mod(source, arena, compile_mode);
+                if (mod == NULL || !_PyAST_Validate(mod)) {
+                    _PyArena_Free(arena);
+                    goto error;
+                }
+                if (_PyCompile_AstOptimize(mod, filename, &cf, optimize,
+                                           arena) < 0) {
+                    _PyArena_Free(arena);
+                    goto error;
+                }
+                result = PyAST_mod2obj(mod);
+            }
+            else {
+                mod_ty mod = PyAST_obj2mod(source, arena, compile_mode);
+                if (mod == NULL || !_PyAST_Validate(mod)) {
+                    _PyArena_Free(arena);
+                    goto error;
+                }
+                result = (PyObject*)_PyAST_Compile(mod, filename,
+                                                   &cf, optimize, arena);
+            }
             _PyArena_Free(arena);
         }
         goto finally;

Python/compile.c

@@ -557,6 +557,24 @@ _PyAST_Compile(mod_ty mod, PyObject *filename, PyCompilerFlags *pflags,
     return co;
 }
 
+int
+_PyCompile_AstOptimize(mod_ty mod, PyObject *filename, PyCompilerFlags *cf,
+                       int optimize, PyArena *arena)
+{
+    PyFutureFeatures future;
+    if (!_PyFuture_FromAST(mod, filename, &future)) {
+        return -1;
+    }
+    int flags = future.ff_features | cf->cf_flags;
+    if (optimize == -1) {
+        optimize = _Py_GetConfig()->optimization_level;
+    }
+    if (!_PyAST_Optimize(mod, arena, optimize, flags)) {
+        return -1;
+    }
+    return 0;
+}
+
 static void
 compiler_free(struct compiler *c)
 {

Python/pythonrun.c

@@ -22,7 +22,6 @@
 #include "pycore_pyerrors.h"      // _PyErr_GetRaisedException, _Py_Offer_Suggestions
 #include "pycore_pylifecycle.h"   // _Py_UnhandledKeyboardInterrupt
 #include "pycore_pystate.h"       // _PyInterpreterState_GET()
-#include "pycore_symtable.h"      // _PyFuture_FromAST()
 #include "pycore_sysmodule.h"     // _PySys_Audit()
 #include "pycore_traceback.h"     // _PyTraceBack_Print_Indented()
 
@@ -1792,24 +1791,6 @@ run_pyc_file(FILE *fp, PyObject *globals, PyObject *locals,
     return NULL;
 }
 
-static int
-ast_optimize(mod_ty mod, PyObject *filename, PyCompilerFlags *cf,
-             int optimize, PyArena *arena)
-{
-    PyFutureFeatures future;
-    if (!_PyFuture_FromAST(mod, filename, &future)) {
-        return -1;
-    }
-    int flags = future.ff_features | cf->cf_flags;
-    if (optimize == -1) {
-        optimize = _Py_GetConfig()->optimization_level;
-    }
-    if (!_PyAST_Optimize(mod, arena, optimize, flags)) {
-        return -1;
-    }
-    return 0;
-}
-
 PyObject *
 Py_CompileStringObject(const char *str, PyObject *filename, int start,
                        PyCompilerFlags *flags, int optimize)
@@ -1827,8 +1808,7 @@ Py_CompileStringObject(const char *str, PyObject *filename, int start,
     }
     if (flags && (flags->cf_flags & PyCF_ONLY_AST)) {
         if ((flags->cf_flags & PyCF_OPTIMIZED_AST) == PyCF_OPTIMIZED_AST) {
-            if (ast_optimize(mod, filename, flags, optimize, arena) < 0) {
-                _PyArena_Free(arena);
+            if (_PyCompile_AstOptimize(mod, filename, flags, optimize, arena) < 0) {
                 return NULL;
             }
         }