diff --git a/Include/internal/pycore_compile.h b/Include/internal/pycore_compile.h index ad657c0f0fc..0167e590e63 100644 --- 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 f3c7229f0b6..68de4d6f3f1 100644 --- 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 ee3ba6ab07b..dbc706ae7b4 100644 --- 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 d06efcf6ba3..787f53fae35 100644 --- 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 3260dba57ea..4b2f70a7ef0 100644 --- 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 a3de7792cf5..05b7dfa61d1 100644 --- 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; } }