test_walker.py 文件源码

def test_walk_field_transform_inspect():

    @ast_walker
    def names_and_incremented_nums(state, node, walk_field, **kwds):
        if isinstance(node, ast.Assign):
            state, value_node = walk_field(state, node.value)
            new_node = replace_fields(node, targets=node.targets, value=value_node)
            new_state = state.update(objs=state.objs.add(node.targets[0].id))
            return new_state, new_node
        elif isinstance(node, ast.Num):
            return state.update(objs=state.objs.add(node.n)), ast.Num(n=node.n + 1)
        else:
            return state, node

    node = get_ast(dummy)
    state, new_node = names_and_incremented_nums(dict(objs=immutableset()), node)
    assert state.objs == set(['a', 'c', 1, 4])
    assert_ast_equal(new_node, get_ast(
        """
        def dummy(x, y):
            c = 5
            a = 2
        """))