mirror of https://github.com/python/cpython
gh-122666: Tests for ast optimizations (#122667)
Co-authored-by: Sergey B Kirpichev <skirpichev@gmail.com> Co-authored-by: Victor Stinner <vstinner@python.org> Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>
This commit is contained in:
parent
1eed0f968f
commit
9f9b00d52c
|
@ -380,7 +380,6 @@ def main():
|
|||
print("]")
|
||||
print("main()")
|
||||
raise SystemExit
|
||||
unittest.main()
|
||||
|
||||
#### EVERYTHING BELOW IS GENERATED BY python Lib/test/test_ast/snippets.py -g #####
|
||||
exec_results = [
|
||||
|
|
|
@ -3035,3 +3035,216 @@ class ASTMainTests(unittest.TestCase):
|
|||
self.assertEqual(expected.splitlines(),
|
||||
res.out.decode("utf8").splitlines())
|
||||
self.assertEqual(res.rc, 0)
|
||||
|
||||
|
||||
class ASTOptimiziationTests(unittest.TestCase):
|
||||
binop = {
|
||||
"+": ast.Add(),
|
||||
"-": ast.Sub(),
|
||||
"*": ast.Mult(),
|
||||
"/": ast.Div(),
|
||||
"%": ast.Mod(),
|
||||
"<<": ast.LShift(),
|
||||
">>": ast.RShift(),
|
||||
"|": ast.BitOr(),
|
||||
"^": ast.BitXor(),
|
||||
"&": ast.BitAnd(),
|
||||
"//": ast.FloorDiv(),
|
||||
"**": ast.Pow(),
|
||||
}
|
||||
|
||||
unaryop = {
|
||||
"~": ast.Invert(),
|
||||
"+": ast.UAdd(),
|
||||
"-": ast.USub(),
|
||||
}
|
||||
|
||||
def wrap_expr(self, expr):
|
||||
return ast.Module(body=[ast.Expr(value=expr)])
|
||||
|
||||
def wrap_for(self, for_statement):
|
||||
return ast.Module(body=[for_statement])
|
||||
|
||||
def assert_ast(self, code, non_optimized_target, optimized_target):
|
||||
non_optimized_tree = ast.parse(code, optimize=-1)
|
||||
optimized_tree = ast.parse(code, optimize=1)
|
||||
|
||||
# Is a non-optimized tree equal to a non-optimized target?
|
||||
self.assertTrue(
|
||||
ast.compare(non_optimized_tree, non_optimized_target),
|
||||
f"{ast.dump(non_optimized_target)} must equal "
|
||||
f"{ast.dump(non_optimized_tree)}",
|
||||
)
|
||||
|
||||
# Is a optimized tree equal to a non-optimized target?
|
||||
self.assertFalse(
|
||||
ast.compare(optimized_tree, non_optimized_target),
|
||||
f"{ast.dump(non_optimized_target)} must not equal "
|
||||
f"{ast.dump(non_optimized_tree)}"
|
||||
)
|
||||
|
||||
# Is a optimized tree is equal to an optimized target?
|
||||
self.assertTrue(
|
||||
ast.compare(optimized_tree, optimized_target),
|
||||
f"{ast.dump(optimized_target)} must equal "
|
||||
f"{ast.dump(optimized_tree)}",
|
||||
)
|
||||
|
||||
def test_folding_binop(self):
|
||||
code = "1 %s 1"
|
||||
operators = self.binop.keys()
|
||||
|
||||
def create_binop(operand, left=ast.Constant(1), right=ast.Constant(1)):
|
||||
return ast.BinOp(left=left, op=self.binop[operand], right=right)
|
||||
|
||||
for op in operators:
|
||||
result_code = code % op
|
||||
non_optimized_target = self.wrap_expr(create_binop(op))
|
||||
optimized_target = self.wrap_expr(ast.Constant(value=eval(result_code)))
|
||||
|
||||
with self.subTest(
|
||||
result_code=result_code,
|
||||
non_optimized_target=non_optimized_target,
|
||||
optimized_target=optimized_target
|
||||
):
|
||||
self.assert_ast(result_code, non_optimized_target, optimized_target)
|
||||
|
||||
# Multiplication of constant tuples must be folded
|
||||
code = "(1,) * 3"
|
||||
non_optimized_target = self.wrap_expr(create_binop("*", ast.Tuple(elts=[ast.Constant(value=1)]), ast.Constant(value=3)))
|
||||
optimized_target = self.wrap_expr(ast.Constant(eval(code)))
|
||||
|
||||
self.assert_ast(code, non_optimized_target, optimized_target)
|
||||
|
||||
def test_folding_unaryop(self):
|
||||
code = "%s1"
|
||||
operators = self.unaryop.keys()
|
||||
|
||||
def create_unaryop(operand):
|
||||
return ast.UnaryOp(op=self.unaryop[operand], operand=ast.Constant(1))
|
||||
|
||||
for op in operators:
|
||||
result_code = code % op
|
||||
non_optimized_target = self.wrap_expr(create_unaryop(op))
|
||||
optimized_target = self.wrap_expr(ast.Constant(eval(result_code)))
|
||||
|
||||
with self.subTest(
|
||||
result_code=result_code,
|
||||
non_optimized_target=non_optimized_target,
|
||||
optimized_target=optimized_target
|
||||
):
|
||||
self.assert_ast(result_code, non_optimized_target, optimized_target)
|
||||
|
||||
def test_folding_not(self):
|
||||
code = "not (1 %s (1,))"
|
||||
operators = {
|
||||
"in": ast.In(),
|
||||
"is": ast.Is(),
|
||||
}
|
||||
opt_operators = {
|
||||
"is": ast.IsNot(),
|
||||
"in": ast.NotIn(),
|
||||
}
|
||||
|
||||
def create_notop(operand):
|
||||
return ast.UnaryOp(op=ast.Not(), operand=ast.Compare(
|
||||
left=ast.Constant(value=1),
|
||||
ops=[operators[operand]],
|
||||
comparators=[ast.Tuple(elts=[ast.Constant(value=1)])]
|
||||
))
|
||||
|
||||
for op in operators.keys():
|
||||
result_code = code % op
|
||||
non_optimized_target = self.wrap_expr(create_notop(op))
|
||||
optimized_target = self.wrap_expr(
|
||||
ast.Compare(left=ast.Constant(1), ops=[opt_operators[op]], comparators=[ast.Constant(value=(1,))])
|
||||
)
|
||||
|
||||
with self.subTest(
|
||||
result_code=result_code,
|
||||
non_optimized_target=non_optimized_target,
|
||||
optimized_target=optimized_target
|
||||
):
|
||||
self.assert_ast(result_code, non_optimized_target, optimized_target)
|
||||
|
||||
def test_folding_format(self):
|
||||
code = "'%s' % (a,)"
|
||||
|
||||
non_optimized_target = self.wrap_expr(
|
||||
ast.BinOp(
|
||||
left=ast.Constant(value="%s"),
|
||||
op=ast.Mod(),
|
||||
right=ast.Tuple(elts=[ast.Name(id='a')]))
|
||||
)
|
||||
optimized_target = self.wrap_expr(
|
||||
ast.JoinedStr(
|
||||
values=[
|
||||
ast.FormattedValue(value=ast.Name(id='a'), conversion=115)
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
self.assert_ast(code, non_optimized_target, optimized_target)
|
||||
|
||||
|
||||
def test_folding_tuple(self):
|
||||
code = "(1,)"
|
||||
|
||||
non_optimized_target = self.wrap_expr(ast.Tuple(elts=[ast.Constant(1)]))
|
||||
optimized_target = self.wrap_expr(ast.Constant(value=(1,)))
|
||||
|
||||
self.assert_ast(code, non_optimized_target, optimized_target)
|
||||
|
||||
def test_folding_comparator(self):
|
||||
code = "1 %s %s1%s"
|
||||
operators = [("in", ast.In()), ("not in", ast.NotIn())]
|
||||
braces = [
|
||||
("[", "]", ast.List, (1,)),
|
||||
("{", "}", ast.Set, frozenset({1})),
|
||||
]
|
||||
for left, right, non_optimized_comparator, optimized_comparator in braces:
|
||||
for op, node in operators:
|
||||
non_optimized_target = self.wrap_expr(ast.Compare(
|
||||
left=ast.Constant(1), ops=[node],
|
||||
comparators=[non_optimized_comparator(elts=[ast.Constant(1)])]
|
||||
))
|
||||
optimized_target = self.wrap_expr(ast.Compare(
|
||||
left=ast.Constant(1), ops=[node],
|
||||
comparators=[ast.Constant(value=optimized_comparator)]
|
||||
))
|
||||
self.assert_ast(code % (op, left, right), non_optimized_target, optimized_target)
|
||||
|
||||
def test_folding_iter(self):
|
||||
code = "for _ in %s1%s: pass"
|
||||
braces = [
|
||||
("[", "]", ast.List, (1,)),
|
||||
("{", "}", ast.Set, frozenset({1})),
|
||||
]
|
||||
|
||||
for left, right, ast_cls, optimized_iter in braces:
|
||||
non_optimized_target = self.wrap_for(ast.For(
|
||||
target=ast.Name(id="_", ctx=ast.Store()),
|
||||
iter=ast_cls(elts=[ast.Constant(1)]),
|
||||
body=[ast.Pass()]
|
||||
))
|
||||
optimized_target = self.wrap_for(ast.For(
|
||||
target=ast.Name(id="_", ctx=ast.Store()),
|
||||
iter=ast.Constant(value=optimized_iter),
|
||||
body=[ast.Pass()]
|
||||
))
|
||||
|
||||
self.assert_ast(code % (left, right), non_optimized_target, optimized_target)
|
||||
|
||||
def test_folding_subscript(self):
|
||||
code = "(1,)[0]"
|
||||
|
||||
non_optimized_target = self.wrap_expr(
|
||||
ast.Subscript(value=ast.Tuple(elts=[ast.Constant(value=1)]), slice=ast.Constant(value=0))
|
||||
)
|
||||
optimized_target = self.wrap_expr(ast.Constant(value=1))
|
||||
|
||||
self.assert_ast(code, non_optimized_target, optimized_target)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
|
Loading…
Reference in New Issue