gh-108113: Make it possible to create an optimized AST (#108154)

This commit is contained in:
Irit Katriel 2023-08-21 17:31:30 +01:00 committed by GitHub
parent 47022a079e
commit 10a91d7e98
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
10 changed files with 124 additions and 14 deletions

View File

@ -2122,10 +2122,12 @@ Async and await
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 @@ and classes for traversing abstract syntax trees:
.. versionchanged:: 3.13
The minimum supported version for feature_version is now (3,7)
The ``optimize`` argument was added.
.. function:: unparse(ast_obj)

View File

@ -85,6 +85,12 @@ Other Language Changes
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 @@ New Modules
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
-----

View File

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

View File

@ -32,13 +32,15 @@ from enum import IntEnum, auto, _simple_enum
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):

View File

@ -357,6 +357,34 @@ class AST_Tests(unittest.TestCase):
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)

View File

@ -369,6 +369,7 @@ class BuiltinTest(unittest.TestCase):
(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))
@ -517,6 +518,28 @@ class BuiltinTest(unittest.TestCase):
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')

View File

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

View File

@ -1208,6 +1208,9 @@ class ASTModuleVisitor(PickleVisitor):
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)

3
Python/Python-ast.c generated
View File

@ -12659,6 +12659,9 @@ astmodule_exec(PyObject *m)
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;
}

View File

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