test_walker.py 文件源码

def test_walker():

    @ast_walker
    def process_numbers(state, node, **kwds):
        if isinstance(node, ast.Num):
            return state.update(numbers=state.numbers.add(node.n)), ast.Num(n=node.n + 1)
        else:
            return state, node

    node = get_ast(dummy)
    state, new_node = process_numbers(dict(numbers=immutableset()), node)

    assert state.numbers == set([1, 4])
    assert_ast_equal(new_node, get_ast("""
        def dummy(x, y):
            c = 5
            a = 2
        """))


# Transformations