抽象语法树如何修改节点

抽象语法树如何修改节点

上次讲了怎么查看抽象语法树，这次讲如何修改抽象语法树。

ast模块有两个类可以修改抽象语法树，分别是NodeVisitor和NodeTransformer。讲一下两者的区别，NodeVisitor可以修改抽象语法树相应节点的内容，NodeTransformer可以彻底修改抽象语法树节点，比如有一个赋值节点，可以使用NodeVisitor做一些小修改，但是不能将赋值节点改成其他节点，但是使用NodeTransformer就可以任意修改节点，将赋值节点改成其他节点完全没有问题，所以NodeVisitor能做的事情NodeTransformer都可以做，并且NodeTransformer可以做的更彻底。

这两个类都使用visit_xxx的方法来遍历每个节点，在遍历节点的过程中就可以批量修改某些内容了，看一下代码：

""" 做一些ast改变代码的试验

    1.ast.NodeVisitor可以改变语法节点的内部内容
    2.ast.NodeTransformer可以更改变节点
    ast.NodeTransformer的作用效果更好,更好控制,也更危险

    the statistics of this file:
    lines(count)    understand_level(h/m/l)    classes(count)    functions(count)    fields(count)
    000000000074    ----------------------m    00000000000002    0000000000000000    ~~~~~~~~~~~~1
"""

import time
import ast

import astunparse

__author__ = '与C同行'
ast_str = \
    """
def add(x, y):
    add_outcome = x + y
    return add_outcome

def list_add(give_list, value):
    pow_value = value ** 2
    give_list.append(pow_value) 
      
store_a = 1

outcome = add(2, 1)
print(f'2+1={outcome}, actual equation:2-1=1')
    """


class AddVisitor(ast.NodeVisitor):
    def visit_BinOp(self, node):
        """ 更改加号为减号

        :param node: 节点
        :return: None
        """
        if isinstance(node.op, ast.Add):
            node.op = ast.Sub()
        self.generic_visit(node)


class DefVisitor(ast.NodeTransformer):
    def visit_FunctionDef(self, node):
        """ 将list_add函数更改为一个赋值语句

        :param node: 函数节点
        :return: 更改后的节点
        """
        if node.name == 'list_add':
            node = ast.Assign(targets=[ast.Name(id='store_b', ctx=ast.Store())],
                              value=ast.Constant(value=2, kind=None))
        self.generic_visit(node)
        return node


if __name__ == '__main__':
    print(f'当前时间:{time.ctime()}')
    ast_parse = ast.parse(ast_str)
    # print(astunparse.dump(ast_parse))
    visitor_change = AddVisitor()
    visitor_change.visit(ast_parse)
    print('改变节点内容:', end='')
    print(astunparse.unparse(ast_parse))
    visitor_transformer = DefVisitor()
    visitor_transformer.visit(ast_parse)
    print('更换节点:', end='')
    print(astunparse.unparse(ast_parse))
    print('改变代码之后执行结果:')
    exec(astunparse.unparse(ast_parse))

上面的代码分别使用了NodeVisitor和NodeTransformer的继承子类来修改节点，NodeVisitor的子类将二元操作符的加号改成减号，NodeTransformer的子类将函数节点改成一个赋值节点，看一下结果：

可以看到修改之后的抽象语法树改成了我们想要的样子，最后执行修改之后的抽象语法树，执行结果显示outcome = add(2, 1)的结果不是3而是1，这说明我们成功地将加号改为了减号。

想了解更多python的知识可以关注我的微信公众号“与C同行”：