def test_assert_statements(self):
optimizer = AstOptimizer(optimize=True)
non_optimizer = AstOptimizer(optimize=False)
code = """def f(a, b): assert a == b, 'lol'"""
tree = ast.parse(code)
optimized = optimizer.visit(tree)
# Function body should contain the assert
self.assertIsInstance(optimized.body[0].body[0], ast.Assert)
unoptimized = non_optimizer.visit(tree)
# Function body should contain the assert
self.assertIsInstance(unoptimized.body[0].body[0], ast.Assert)
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_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)