X Tutup
Skip to content

Commit 2dfbd4f

Browse files
iritkatrieliritkatriel
authored
pythongh-108113: Make it possible to optimize an AST (python#108282)
1 parent 79fdacc commit 2dfbd4f

File tree

6 files changed

+77
-48
lines changed

6 files changed

+77
-48
lines changed

‎Include/internal/pycore_compile.h‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,14 @@ PyAPI_FUNC(PyCodeObject*) _PyAST_Compile(
1919
int optimize,
2020
struct _arena *arena);
2121

22+
/* AST optimizations */
23+
PyAPI_FUNC(int) _PyCompile_AstOptimize(
24+
struct _mod *mod,
25+
PyObject *filename,
26+
PyCompilerFlags *flags,
27+
int optimize,
28+
struct _arena *arena);
29+
2230
static const _PyCompilerSrcLocation NO_LOCATION = {-1, -1, -1, -1};
2331

2432
extern int _PyAST_Optimize(

‎Lib/test/test_ast.py‎

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -361,14 +361,16 @@ def test_optimization_levels__debug__(self):
361361
cases = [(-1, '__debug__'), (0, '__debug__'), (1, False), (2, False)]
362362
for (optval, expected) in cases:
363363
with self.subTest(optval=optval, expected=expected):
364-
res = ast.parse("__debug__", optimize=optval)
365-
self.assertIsInstance(res.body[0], ast.Expr)
366-
if isinstance(expected, bool):
367-
self.assertIsInstance(res.body[0].value, ast.Constant)
368-
self.assertEqual(res.body[0].value.value, expected)
369-
else:
370-
self.assertIsInstance(res.body[0].value, ast.Name)
371-
self.assertEqual(res.body[0].value.id, expected)
364+
res1 = ast.parse("__debug__", optimize=optval)
365+
res2 = ast.parse(ast.parse("__debug__"), optimize=optval)
366+
for res in [res1, res2]:
367+
self.assertIsInstance(res.body[0], ast.Expr)
368+
if isinstance(expected, bool):
369+
self.assertIsInstance(res.body[0].value, ast.Constant)
370+
self.assertEqual(res.body[0].value.value, expected)
371+
else:
372+
self.assertIsInstance(res.body[0].value, ast.Name)
373+
self.assertEqual(res.body[0].value.id, expected)
372374

373375
def test_optimization_levels_const_folding(self):
374376
folded = ('Expr', (1, 0, 1, 5), ('Constant', (1, 0, 1, 5), 3, None))
@@ -381,9 +383,11 @@ def test_optimization_levels_const_folding(self):
381383
cases = [(-1, not_folded), (0, not_folded), (1, folded), (2, folded)]
382384
for (optval, expected) in cases:
383385
with self.subTest(optval=optval):
384-
tree = ast.parse("1 + 2", optimize=optval)
385-
res = to_tuple(tree.body[0])
386-
self.assertEqual(res, expected)
386+
tree1 = ast.parse("1 + 2", optimize=optval)
387+
tree2 = ast.parse(ast.parse("1 + 2"), optimize=optval)
388+
for tree in [tree1, tree2]:
389+
res = to_tuple(tree.body[0])
390+
self.assertEqual(res, expected)
387391

388392
def test_invalid_position_information(self):
389393
invalid_linenos = [

‎Lib/test/test_builtin.py‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -521,9 +521,10 @@ def test_compile_async_generator(self):
521521
def test_compile_ast(self):
522522
args = ("a*(1+2)", "f.py", "exec")
523523
raw = compile(*args, flags = ast.PyCF_ONLY_AST).body[0]
524-
opt = compile(*args, flags = ast.PyCF_OPTIMIZED_AST).body[0]
524+
opt1 = compile(*args, flags = ast.PyCF_OPTIMIZED_AST).body[0]
525+
opt2 = compile(ast.parse(args[0]), *args[1:], flags = ast.PyCF_OPTIMIZED_AST).body[0]
525526

526-
for tree in (raw, opt):
527+
for tree in (raw, opt1, opt2):
527528
self.assertIsInstance(tree.value, ast.BinOp)
528529
self.assertIsInstance(tree.value.op, ast.Mult)
529530
self.assertIsInstance(tree.value.left, ast.Name)
@@ -536,9 +537,10 @@ def test_compile_ast(self):
536537
self.assertIsInstance(raw_right.right, ast.Constant)
537538
self.assertEqual(raw_right.right.value, 2)
538539

539-
opt_right = opt.value.right # expect Constant(3)
540-
self.assertIsInstance(opt_right, ast.Constant)
541-
self.assertEqual(opt_right.value, 3)
540+
for opt in [opt1, opt2]:
541+
opt_right = opt.value.right # expect Constant(3)
542+
self.assertIsInstance(opt_right, ast.Constant)
543+
self.assertEqual(opt_right.value, 3)
542544

543545
def test_delattr(self):
544546
sys.spam = 1

‎Python/bltinmodule.c‎

Lines changed: 28 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -804,23 +804,40 @@ builtin_compile_impl(PyObject *module, PyObject *source, PyObject *filename,
804804
if (is_ast == -1)
805805
goto error;
806806
if (is_ast) {
807-
if (flags & PyCF_ONLY_AST) {
807+
if ((flags & PyCF_OPTIMIZED_AST) == PyCF_ONLY_AST) {
808+
// return an un-optimized AST
808809
result = Py_NewRef(source);
809810
}
810811
else {
811-
PyArena *arena;
812-
mod_ty mod;
812+
// Return an optimized AST or code object
813813

814-
arena = _PyArena_New();
815-
if (arena == NULL)
816-
goto error;
817-
mod = PyAST_obj2mod(source, arena, compile_mode);
818-
if (mod == NULL || !_PyAST_Validate(mod)) {
819-
_PyArena_Free(arena);
814+
PyArena *arena = _PyArena_New();
815+
if (arena == NULL) {
820816
goto error;
821817
}
822-
result = (PyObject*)_PyAST_Compile(mod, filename,
823-
&cf, optimize, arena);
818+
819+
if (flags & PyCF_ONLY_AST) {
820+
mod_ty mod = PyAST_obj2mod(source, arena, compile_mode);
821+
if (mod == NULL || !_PyAST_Validate(mod)) {
822+
_PyArena_Free(arena);
823+
goto error;
824+
}
825+
if (_PyCompile_AstOptimize(mod, filename, &cf, optimize,
826+
arena) < 0) {
827+
_PyArena_Free(arena);
828+
goto error;
829+
}
830+
result = PyAST_mod2obj(mod);
831+
}
832+
else {
833+
mod_ty mod = PyAST_obj2mod(source, arena, compile_mode);
834+
if (mod == NULL || !_PyAST_Validate(mod)) {
835+
_PyArena_Free(arena);
836+
goto error;
837+
}
838+
result = (PyObject*)_PyAST_Compile(mod, filename,
839+
&cf, optimize, arena);
840+
}
824841
_PyArena_Free(arena);
825842
}
826843
goto finally;

‎Python/compile.c‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -557,6 +557,24 @@ _PyAST_Compile(mod_ty mod, PyObject *filename, PyCompilerFlags *pflags,
557557
return co;
558558
}
559559

560+
int
561+
_PyCompile_AstOptimize(mod_ty mod, PyObject *filename, PyCompilerFlags *cf,
562+
int optimize, PyArena *arena)
563+
{
564+
PyFutureFeatures future;
565+
if (!_PyFuture_FromAST(mod, filename, &future)) {
566+
return -1;
567+
}
568+
int flags = future.ff_features | cf->cf_flags;
569+
if (optimize == -1) {
570+
optimize = _Py_GetConfig()->optimization_level;
571+
}
572+
if (!_PyAST_Optimize(mod, arena, optimize, flags)) {
573+
return -1;
574+
}
575+
return 0;
576+
}
577+
560578
static void
561579
compiler_free(struct compiler *c)
562580
{

‎Python/pythonrun.c‎

Lines changed: 1 addition & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@
2222
#include "pycore_pyerrors.h" // _PyErr_GetRaisedException, _Py_Offer_Suggestions
2323
#include "pycore_pylifecycle.h" // _Py_UnhandledKeyboardInterrupt
2424
#include "pycore_pystate.h" // _PyInterpreterState_GET()
25-
#include "pycore_symtable.h" // _PyFuture_FromAST()
2625
#include "pycore_sysmodule.h" // _PySys_Audit()
2726
#include "pycore_traceback.h" // _PyTraceBack_Print_Indented()
2827

@@ -1792,24 +1791,6 @@ run_pyc_file(FILE *fp, PyObject *globals, PyObject *locals,
17921791
return NULL;
17931792
}
17941793

1795-
static int
1796-
ast_optimize(mod_ty mod, PyObject *filename, PyCompilerFlags *cf,
1797-
int optimize, PyArena *arena)
1798-
{
1799-
PyFutureFeatures future;
1800-
if (!_PyFuture_FromAST(mod, filename, &future)) {
1801-
return -1;
1802-
}
1803-
int flags = future.ff_features | cf->cf_flags;
1804-
if (optimize == -1) {
1805-
optimize = _Py_GetConfig()->optimization_level;
1806-
}
1807-
if (!_PyAST_Optimize(mod, arena, optimize, flags)) {
1808-
return -1;
1809-
}
1810-
return 0;
1811-
}
1812-
18131794
PyObject *
18141795
Py_CompileStringObject(const char *str, PyObject *filename, int start,
18151796
PyCompilerFlags *flags, int optimize)
@@ -1827,8 +1808,7 @@ Py_CompileStringObject(const char *str, PyObject *filename, int start,
18271808
}
18281809
if (flags && (flags->cf_flags & PyCF_ONLY_AST)) {
18291810
if ((flags->cf_flags & PyCF_OPTIMIZED_AST) == PyCF_OPTIMIZED_AST) {
1830-
if (ast_optimize(mod, filename, flags, optimize, arena) < 0) {
1831-
_PyArena_Free(arena);
1811+
if (_PyCompile_AstOptimize(mod, filename, flags, optimize, arena) < 0) {
18321812
return NULL;
18331813
}
18341814
}

0 commit comments

Comments
 (0)
X Tutup