def test_fold_nonconst_list_to_tuple_in_comparisons(self):
optimizer = AstOptimizer()
tree = ast.parse("[a for a in b if a.c in [e, f]]")
optimized = optimizer.visit(tree)
self.assertEqual(
to_expr(
optimized.body[0].value.generators[0].ifs[0].comparators[0]),
"(e, f)",
)
def test_rewrite_node(self):
class TestRewriter(ASTRewriter):
def visitName(self, node):
if isinstance(node, ast.Name) and node.id == "z":
return ast.Name("foo", ast.Load())
return node
tree = ast.parse("x + y + z").body[0].value
rewriter = TestRewriter()
new_tree = rewriter.visit(tree)
self.assertIsNotNone(new_tree)
self.assertNotEqual(new_tree, tree, "tree should be rewritten")
self.assertEqual(to_expr(new_tree), "x + y + foo")
self.assertEqual(tree.left, new_tree.left,
"Unchanged nodes should be the same")
def test_simple_examples(self):
examples = [
"a",
"+a",
"-a",
"not a",
"~a",
"x is y",
"x is not y",
"x < y",
"x > y",
"x <= y",
"x >= y",
"x == y",
"x != y",
"None",
"True",
"False",
"...",
"42.0",
"42j",
"'abc'",
"b'abc'",
"foo.bar",
"42 .bar",
"42.0.bar" if sys.version_info >= (3, 8) else "42.0 .bar",
"42j.bar" if sys.version_info >= (3, 8) else "42j .bar",
"()",
"(1,)",
"(1, 2)",
"[]",
"[1, 2]",
"{1, 2}",
"{}",
"{1: 2}",
"{1: 2, 3: 4}",
"a + 2",
"a - 2",
"a * 2",
"a @ 2",
"a / 2",
"a % 2",
"a << 2",
"a >> 2",
"a | 2",
"a ^ 2",
"a // 2",
"a ** 2",
"a[b]",
"a[b:c]",
"a[:c]",
"a[b:]",
"a[:]",
"a[a:b:c]",
"a[::c]",
"a[b,]",
"a[b, c:d]",
"a(b)",
"a(b, c)",
"a(b, *c)",
"a(b, *c, **foo)",
"a(b=1)",
"a(b=1, **foo)",
"a(a, b=1, **foo)",
"a(a, b=1, **foo)",
"(yield)",
"(yield 42)",
"(yield from [])",
"lambda: 42",
"lambda a: 42",
"lambda a=2: 42",
"lambda*foo: 42", # oddity of how CPython unparses this...
"lambda a, *foo: 42",
"lambda a, *, b=2: 42",
"lambda*, b: 42", # oddity of how CPython unparses this...
"lambda*, b=2: 42", # oddity of how CPython unparses this...
"lambda x: (x, x)",
"1 if True else 2",
"f'foo'",
"f'foo{bar}'",
"f'foo{bar!a}'",
# joined strings get combined
("f'foo{bar:N}'f'foo{bar:N}'", "f'foo{bar:N}foo{bar:N}'"),
"f'foo{ {2: 3}}'",
"f'{(lambda x: x)}'",
"[x for x in abc]",
"{x for x in abc}",
"{x: y for x, y in abc}",
"[x for x in abc if x]",
"{x for x in abc if x}",
"{x: y for x, y in abc if x}",
"[x for x in abc for z in foo]",
"{x for x in abc for z in foo}",
"{x: y for x, y in abc for z in foo}",
"[*abc]",
"(x for x in y)",
]
for example in examples:
if isinstance(example, tuple):
example, expected = example
else:
expected = example
tree = ast.parse(example)
self.assertEqual(expected, to_expr(tree.body[0].value))
if sys.version_info >= (3, 7):
# Make sure we match CPython's compilation too
l = {}
exec(
f"from __future__ import annotations\ndef f() -> {example}:\n pass",
l,
l,
)
f = l["f"]
self.assertEqual(expected, f.__annotations__["return"])
def test_ast_optimizer_for(self):
optimizer = AstOptimizer()
tree = ast.parse("for x in [1,2,3]: pass")
optimized = optimizer.visit(tree).body[0]
self.assertEqual(to_expr(optimized.iter), "(1, 2, 3)")
def test_ast_optimizer(self):
cases = [
("+1", "1"),
("--1", "1"),
("~1", "-2"),
("not 1", "False"),
("not x is y", "x is not y"),
("not x is not y", "x is y"),
("not x in y", "x not in y"),
("~1.1", "~1.1"),
("+'str'", "+'str'"),
("1 + 2", "3"),
("1 + 3", "4"),
("'abc' + 'def'", "'abcdef'"),
("b'abc' + b'def'", "b'abcdef'"),
("b'abc' + 'def'", "b'abc' + 'def'"),
("b'abc' + --2", "b'abc' + 2"),
("--2 + 'abc'", "2 + 'abc'"),
("5 - 3", "2"),
("6 - 3", "3"),
("2 * 2", "4"),
("2 * 3", "6"),
("'abc' * 2", "'abcabc'"),
("b'abc' * 2", "b'abcabc'"),
("1 / 2", "0.5"),
("6 / 2", "3.0"),
("6 // 2", "3"),
("5 // 2", "2"),
("2 >> 1", "1"),
("6 >> 1", "3"),
("1 | 2", "3"),
("1 | 1", "1"),
("1 ^ 3", "2"),
("1 ^ 1", "0"),
("1 & 2", "0"),
("1 & 3", "1"),
("'abc' + 1", "'abc' + 1"),
("1 / 0", "1 / 0"),
("1 + None", "1 + None"),
("True + None", "True + None"),
("True + 1", "2"),
("(1, 2)", "(1, 2)"),
("(1, 2) * 2", "(1, 2, 1, 2)"),
("(1, --2, abc)", "(1, 2, abc)"),
("(1, 2)[0]", "1"),
("1[0]", "1[0]"),
("x[+1]", "x[1]"),
("(+1)[x]", "1[x]"),
("[x for x in [1,2,3]]", "[x for x in (1, 2, 3)]"),
("(x for x in [1,2,3])", "(x for x in (1, 2, 3))"),
("{x for x in [1,2,3]}", "{x for x in (1, 2, 3)}"),
("{x for x in [--1,2,3]}", "{x for x in (1, 2, 3)}"),
("{--1 for x in [1,2,3]}", "{1 for x in (1, 2, 3)}"),
("x in [1,2,3]", "x in (1, 2, 3)"),
("x in x in [1,2,3]", "x in x in (1, 2, 3)"),
("x in [1,2,3] in x", "x in [1, 2, 3] in x"),
]
for inp, expected in cases:
optimizer = AstOptimizer()
tree = ast.parse(inp)
optimized = to_expr(optimizer.visit(tree).body[0].value)
self.assertEqual(expected, optimized, "Input was: " + inp)