gh-108113: Make it possible to create an optimized AST (#108154)
diff --git a/Doc/library/ast.rst b/Doc/library/ast.rst
index cd657ae..2237a07 100644
--- a/Doc/library/ast.rst
+++ b/Doc/library/ast.rst
@@ -2122,10 +2122,12 @@
Apart from the node classes, the :mod:`ast` module defines these utility functions
and classes for traversing abstract syntax trees:
-.. function:: parse(source, filename='<unknown>', mode='exec', *, type_comments=False, feature_version=None)
+.. function:: parse(source, filename='<unknown>', mode='exec', *, type_comments=False, feature_version=None, optimize=-1)
Parse the source into an AST node. Equivalent to ``compile(source,
- filename, mode, ast.PyCF_ONLY_AST)``.
+ filename, mode, flags=FLAGS_VALUE, optimize=optimize)``,
+ where ``FLAGS_VALUE`` is ``ast.PyCF_ONLY_AST`` if ``optimize <= 0``
+ and ``ast.PyCF_OPTIMIZED_AST`` otherwise.
If ``type_comments=True`` is given, the parser is modified to check
and return type comments as specified by :pep:`484` and :pep:`526`.
@@ -2171,6 +2173,7 @@
.. versionchanged:: 3.13
The minimum supported version for feature_version is now (3,7)
+ The ``optimize`` argument was added.
.. function:: unparse(ast_obj)
diff --git a/Doc/whatsnew/3.13.rst b/Doc/whatsnew/3.13.rst
index 47b868b..bfab868 100644
--- a/Doc/whatsnew/3.13.rst
+++ b/Doc/whatsnew/3.13.rst
@@ -85,6 +85,12 @@
This change will affect tools using docstrings, like :mod:`doctest`.
(Contributed by Inada Naoki in :gh:`81283`.)
+* The :func:`compile` built-in can now accept a new flag,
+ ``ast.PyCF_OPTIMIZED_AST``, which is similar to ``ast.PyCF_ONLY_AST``
+ except that the returned ``AST`` is optimized according to the value
+ of the ``optimize`` argument.
+ (Contributed by Irit Katriel in :gh:`108113`).
+
New Modules
===========
@@ -94,6 +100,14 @@
Improved Modules
================
+ast
+---
+
+* :func:`ast.parse` now accepts an optional argument ``optimize``
+ which is passed on to the :func:`compile` built-in. This makes it
+ possible to obtain an optimized ``AST``.
+ (Contributed by Irit Katriel in :gh:`108113`).
+
array
-----
diff --git a/Include/cpython/compile.h b/Include/cpython/compile.h
index e6cd39a..ae17cef 100644
--- a/Include/cpython/compile.h
+++ b/Include/cpython/compile.h
@@ -19,9 +19,10 @@
#define PyCF_TYPE_COMMENTS 0x1000
#define PyCF_ALLOW_TOP_LEVEL_AWAIT 0x2000
#define PyCF_ALLOW_INCOMPLETE_INPUT 0x4000
+#define PyCF_OPTIMIZED_AST (0x8000 | PyCF_ONLY_AST)
#define PyCF_COMPILE_MASK (PyCF_ONLY_AST | PyCF_ALLOW_TOP_LEVEL_AWAIT | \
PyCF_TYPE_COMMENTS | PyCF_DONT_IMPLY_DEDENT | \
- PyCF_ALLOW_INCOMPLETE_INPUT)
+ PyCF_ALLOW_INCOMPLETE_INPUT | PyCF_OPTIMIZED_AST)
typedef struct {
int cf_flags; /* bitmask of CO_xxx flags relevant to future */
diff --git a/Lib/ast.py b/Lib/ast.py
index a307f3e..45b9596 100644
--- a/Lib/ast.py
+++ b/Lib/ast.py
@@ -32,13 +32,15 @@
def parse(source, filename='<unknown>', mode='exec', *,
- type_comments=False, feature_version=None):
+ type_comments=False, feature_version=None, optimize=-1):
"""
Parse the source into an AST node.
Equivalent to compile(source, filename, mode, PyCF_ONLY_AST).
Pass type_comments=True to get back type comments where the syntax allows.
"""
flags = PyCF_ONLY_AST
+ if optimize > 0:
+ flags |= PyCF_OPTIMIZED_AST
if type_comments:
flags |= PyCF_TYPE_COMMENTS
if feature_version is None:
@@ -50,7 +52,7 @@ def parse(source, filename='<unknown>', mode='exec', *,
feature_version = minor
# Else it should be an int giving the minor version for 3.x.
return compile(source, filename, mode, flags,
- _feature_version=feature_version)
+ _feature_version=feature_version, optimize=optimize)
def literal_eval(node_or_string):
diff --git a/Lib/test/test_ast.py b/Lib/test/test_ast.py
index 5346b39..f3c7229 100644
--- a/Lib/test/test_ast.py
+++ b/Lib/test/test_ast.py
@@ -357,6 +357,34 @@ def test_ast_validation(self):
tree = ast.parse(snippet)
compile(tree, '<string>', 'exec')
+ 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)
+
+ def test_optimization_levels_const_folding(self):
+ folded = ('Expr', (1, 0, 1, 5), ('Constant', (1, 0, 1, 5), 3, None))
+ not_folded = ('Expr', (1, 0, 1, 5),
+ ('BinOp', (1, 0, 1, 5),
+ ('Constant', (1, 0, 1, 1), 1, None),
+ ('Add',),
+ ('Constant', (1, 4, 1, 5), 2, None)))
+
+ 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)
+
def test_invalid_position_information(self):
invalid_linenos = [
(10, 1), (-10, -11), (10, -11), (-5, -2), (-5, 1)
diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
index f5a5c03..ee3ba6a 100644
--- a/Lib/test/test_builtin.py
+++ b/Lib/test/test_builtin.py
@@ -369,16 +369,17 @@ def f(): """doc"""
(1, False, 'doc', False, False),
(2, False, None, False, False)]
for optval, *expected in values:
+ with self.subTest(optval=optval):
# test both direct compilation and compilation via AST
- codeobjs = []
- codeobjs.append(compile(codestr, "<test>", "exec", optimize=optval))
- tree = ast.parse(codestr)
- codeobjs.append(compile(tree, "<test>", "exec", optimize=optval))
- for code in codeobjs:
- ns = {}
- exec(code, ns)
- rv = ns['f']()
- self.assertEqual(rv, tuple(expected))
+ codeobjs = []
+ codeobjs.append(compile(codestr, "<test>", "exec", optimize=optval))
+ tree = ast.parse(codestr)
+ codeobjs.append(compile(tree, "<test>", "exec", optimize=optval))
+ for code in codeobjs:
+ ns = {}
+ exec(code, ns)
+ rv = ns['f']()
+ self.assertEqual(rv, tuple(expected))
def test_compile_top_level_await_no_coro(self):
"""Make sure top level non-await codes get the correct coroutine flags"""
@@ -517,6 +518,28 @@ def test_compile_async_generator(self):
exec(co, glob)
self.assertEqual(type(glob['ticker']()), AsyncGeneratorType)
+ 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]
+
+ for tree in (raw, opt):
+ self.assertIsInstance(tree.value, ast.BinOp)
+ self.assertIsInstance(tree.value.op, ast.Mult)
+ self.assertIsInstance(tree.value.left, ast.Name)
+ self.assertEqual(tree.value.left.id, 'a')
+
+ raw_right = raw.value.right # expect BinOp(1, '+', 2)
+ self.assertIsInstance(raw_right, ast.BinOp)
+ self.assertIsInstance(raw_right.left, ast.Constant)
+ self.assertEqual(raw_right.left.value, 1)
+ 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)
+
def test_delattr(self):
sys.spam = 1
delattr(sys, 'spam')
diff --git a/Misc/NEWS.d/next/Core and Builtins/2023-08-18-18-21-27.gh-issue-108113.1h0poE.rst b/Misc/NEWS.d/next/Core and Builtins/2023-08-18-18-21-27.gh-issue-108113.1h0poE.rst
new file mode 100644
index 0000000..6668057
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2023-08-18-18-21-27.gh-issue-108113.1h0poE.rst
@@ -0,0 +1,8 @@
+The :func:`compile` built-in can now accept a new flag,
+``ast.PyCF_OPTIMIZED_AST``, which is similar to ``ast.PyCF_ONLY_AST``
+except that the returned ``AST`` is optimized according to the value
+of the ``optimize`` argument.
+
+:func:`ast.parse` now accepts an optional argument ``optimize``
+which is passed on to the :func:`compile` built-in. This makes it
+possible to obtain an optimized ``AST``.
diff --git a/Parser/asdl_c.py b/Parser/asdl_c.py
index 2a36610..1733cd4 100755
--- a/Parser/asdl_c.py
+++ b/Parser/asdl_c.py
@@ -1208,6 +1208,9 @@ def visitModule(self, mod):
self.emit('if (PyModule_AddIntMacro(m, PyCF_TYPE_COMMENTS) < 0) {', 1)
self.emit("return -1;", 2)
self.emit('}', 1)
+ self.emit('if (PyModule_AddIntMacro(m, PyCF_OPTIMIZED_AST) < 0) {', 1)
+ self.emit("return -1;", 2)
+ self.emit('}', 1)
for dfn in mod.dfns:
self.visit(dfn)
self.emit("return 0;", 1)
diff --git a/Python/Python-ast.c b/Python/Python-ast.c
index 8047b12..60dd121 100644
--- a/Python/Python-ast.c
+++ b/Python/Python-ast.c
@@ -12659,6 +12659,9 @@
if (PyModule_AddIntMacro(m, PyCF_TYPE_COMMENTS) < 0) {
return -1;
}
+ if (PyModule_AddIntMacro(m, PyCF_OPTIMIZED_AST) < 0) {
+ return -1;
+ }
if (PyModule_AddObjectRef(m, "mod", state->mod_type) < 0) {
return -1;
}
diff --git a/Python/pythonrun.c b/Python/pythonrun.c
index 721c527..b2e04cf 100644
--- a/Python/pythonrun.c
+++ b/Python/pythonrun.c
@@ -21,6 +21,7 @@
#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()
@@ -1790,6 +1791,24 @@ 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)
@@ -1806,6 +1825,12 @@ Py_CompileStringObject(const char *str, PyObject *filename, int start,
return NULL;
}
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);
+ return NULL;
+ }
+ }
PyObject *result = PyAST_mod2obj(mod);
_PyArena_Free(arena);
return result;
