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

This commit is contained in:
Irit Katriel 2023-08-23 09:01:17 +01:00 committed by GitHub
parent 79fdacc005
commit 2dfbd4f36d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 77 additions and 48 deletions

View File

@ -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(

View File

@ -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 = [

View File

@ -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

View File

@ -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;

View File

@ -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)
{

View File

@ -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;
}
}