def test_basic_unroll():
tree = ast.parse("""
def foo():
for i in ast_tools.macros.unroll(range(8)):
print(i)
""")
assert astor.to_source(unroll_for_loops(tree, globals())) == """\
def test_basic_inside_if():
src = """\
def foo(x):
if x:
for i in ast_tools.macros.unroll(range(8)):
print(i)
return x + 1 if x % 2 else x
else:
print(x)
for j in ast_tools.macros.unroll(range(2)):
print(j - 1)
"""
unrolled_src = """\
def foo(x):
if x:
print(0)
print(1)
print(2)
print(3)
print(4)
print(5)
print(6)
print(7)
return x + 1 if x % 2 else x
else:
print(x)
print(0 - 1)
print(1 - 1)
"""
tree = cst.parse_module(src)
unrolled_tree = unroll_for_loops(tree, globals())
assert to_module(unrolled_tree).code == unrolled_src
def test_basic_env():
tree = ast.parse("""
def foo(x):
for i in ast_tools.macros.unroll(range(j)):
print(i)
""")
env = dict(globals(), **{"j": 2})
assert astor.to_source(unroll_for_loops(tree, env)) == """\
def test_basic_inside_if():
tree = ast.parse("""
def foo(x):
if x:
for i in ast_tools.macros.unroll(range(8)):
print(i)
return x + 1 if x % 2 else x
else:
print(x)
for j in ast_tools.macros.unroll(range(2)):
print(j - 1)
""")
assert astor.to_source(unroll_for_loops(tree, globals())) == """\
def test_basic_env():
src = """\
def foo(x):
for i in ast_tools.macros.unroll(range(j)):
print(i)
"""
unrolled_src = """\
def foo(x):
print(0)
print(1)
"""
env = dict(globals(), **{"j": 2})
tree = cst.parse_module(src)
unrolled_tree = unroll_for_loops(tree, env)
assert to_module(unrolled_tree).code == unrolled_src
def test_basic_unroll():
src = """\
def foo():
for i in ast_tools.macros.unroll(range(8)):
print(i)
"""
unrolled_src = """\
def foo():
print(0)
print(1)
print(2)
print(3)
print(4)
print(5)
print(6)
print(7)
"""
tree = cst.parse_module(src)
unrolled_tree = unroll_for_loops(tree, globals())
assert to_module(unrolled_tree).code == unrolled_src
def rewrite(self, tree: cst.CSTNode, env: SymbolTable,
metadata: tp.MutableMapping) -> PASS_ARGS_T:
return unroll_for_loops(tree, env), env, metadata