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