抽象语法树如何修改节点
上次讲了怎么查看抽象语法树,这次讲如何修改抽象语法树。
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同行”: