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

diff --git a/Include/internal/pycore_compile.h b/Include/internal/pycore_compile.h
index ad657c0f0fcedca8d86b9a4f8ba525c0983f9451..0167e590e63fb1db2e8dae367f21304e867fbf64 100644 (file)
--- a/Include/internal/pycore_compile.h
+++ b/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(

diff --git a/Lib/test/test_ast.py b/Lib/test/test_ast.py
index f3c7229f0b6c76a89d2b74681a91c63578070963..68de4d6f3f1923609f11da4591d104aa7a125a61 100644 (file)
--- a/Lib/test/test_ast.py
+++ b/Lib/test/test_ast.py
@@ -361,14 +361,16 @@ class AST_Tests(unittest.TestCase):
         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 @@ class AST_Tests(unittest.TestCase):
         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 = [

diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
index ee3ba6ab07bbdfe1a1ea7f2282283d7e8fe9cd08..dbc706ae7b41f57c41b41a213e0f5348efaef52b 100644 (file)
--- a/Lib/test/test_builtin.py
+++ b/Lib/test/test_builtin.py
@@ -521,9 +521,10 @@ class BuiltinTest(unittest.TestCase):
     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 @@ class BuiltinTest(unittest.TestCase):
         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

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index d06efcf6ba3687a03d84b0e3d638325b927f9b0a..787f53fae354dbc4f2b907fa9ac14c45fbe33ef2 100644 (file)
--- a/Python/bltinmodule.c
+++ b/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;

diff --git a/Python/compile.c b/Python/compile.c
index 3260dba57eac8f97975438811fa08eed610b100c..4b2f70a7ef01d779f9dac8ac287ab1fa91fd484d 100644 (file)
--- a/Python/compile.c
+++ b/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)
 {

diff --git a/Python/pythonrun.c b/Python/pythonrun.c
index a3de7792cf5bc5909f8a7a159e746220c10733c5..05b7dfa61d1a52124e83a64469f36d81723f17f3 100644 (file)
--- a/Python/pythonrun.c
+++ b/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 @@ error:
     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;
             }
         }