From 10a91d7e98d847b05292eab828ff9ae51308d3ee Mon Sep 17 00:00:00 2001 From: Irit Katriel <1055913+iritkatriel@users.noreply.github.com> Date: Mon, 21 Aug 2023 17:31:30 +0100 Subject: [PATCH] gh-108113: Make it possible to create an optimized AST (#108154) --- Doc/library/ast.rst | 7 +++- Doc/whatsnew/3.13.rst | 14 +++++++ Include/cpython/compile.h | 3 +- Lib/ast.py | 6 ++- Lib/test/test_ast.py | 28 +++++++++++++ Lib/test/test_builtin.py | 41 +++++++++++++++---- ...-08-18-18-21-27.gh-issue-108113.1h0poE.rst | 8 ++++ Parser/asdl_c.py | 3 ++ Python/Python-ast.c | 3 ++ Python/pythonrun.c | 25 +++++++++++ 10 files changed, 124 insertions(+), 14 deletions(-) create mode 100644 Misc/NEWS.d/next/Core and Builtins/2023-08-18-18-21-27.gh-issue-108113.1h0poE.rst diff --git a/Doc/library/ast.rst b/Doc/library/ast.rst index cd657aedf6d..2237a07eb9d 100644 --- a/Doc/library/ast.rst +++ b/Doc/library/ast.rst @@ -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='', mode='exec', *, type_comments=False, feature_version=None) +.. function:: parse(source, filename='', 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) diff --git a/Doc/whatsnew/3.13.rst b/Doc/whatsnew/3.13.rst index 47b868bad31..bfab868d1c5 100644 --- a/Doc/whatsnew/3.13.rst +++ b/Doc/whatsnew/3.13.rst @@ -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 ----- diff --git a/Include/cpython/compile.h b/Include/cpython/compile.h index e6cd39af2ba..ae17cef554f 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 a307f3ecd06..45b95963f81 100644 --- a/Lib/ast.py +++ b/Lib/ast.py @@ -32,13 +32,15 @@ from enum import IntEnum, auto, _simple_enum def parse(source, filename='', 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='', 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 5346b39043f..f3c7229f0b6 100644 --- a/Lib/test/test_ast.py +++ b/Lib/test/test_ast.py @@ -357,6 +357,34 @@ class AST_Tests(unittest.TestCase): tree = ast.parse(snippet) compile(tree, '', '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 f5a5c037f1b..ee3ba6ab07b 100644 --- a/Lib/test/test_builtin.py +++ b/Lib/test/test_builtin.py @@ -369,16 +369,17 @@ 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, "", "exec", optimize=optval)) - tree = ast.parse(codestr) - codeobjs.append(compile(tree, "", "exec", optimize=optval)) - for code in codeobjs: - ns = {} - exec(code, ns) - rv = ns['f']() - self.assertEqual(rv, tuple(expected)) + codeobjs = [] + codeobjs.append(compile(codestr, "", "exec", optimize=optval)) + tree = ast.parse(codestr) + codeobjs.append(compile(tree, "", "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 @@ 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') 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 00000000000..66680578c9b --- /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 2a36610527f..1733cd4b15a 100755 --- a/Parser/asdl_c.py +++ b/Parser/asdl_c.py @@ -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) diff --git a/Python/Python-ast.c b/Python/Python-ast.c index 8047b1259c5..60dd121d60b 100644 --- a/Python/Python-ast.c +++ b/Python/Python-ast.c @@ -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; } diff --git a/Python/pythonrun.c b/Python/pythonrun.c index 721c527745c..b2e04cfa317 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 @@ 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;