python中的ast模块介绍和使用

ast模块

官方文档：ast — Abstract Syntax Trees

教程文档：Getting to and from ASTs

ast模块简介

参考文章：python compiler.ast_Python Ast介绍及应用

Python官方提供的CPython解释器对python源码的处理过程如下：

Parse source code into a parse tree (Parser/pgen.c)

Transform parse tree into an Abstract Syntax Tree (Python/ast.c)

Transform AST into a Control Flow Graph (Python/compile.c)

Emit bytecode based on the Control Flow Graph (Python/compile.c)

即实际python代码的处理过程如下：

源代码解析 --> 语法树 --> 抽象语法树(AST) --> 控制流程图 --> 字节码

上述过程在python2.5之后被应用。python源码首先被解析成语法树，随后又转换成抽象语法树。在抽象语法树中我们可以看到python源码文件中的语法结构。

查看ast抽象语法树

import ast
root_node = ast.parse("print('hello world')")

print(root_node)
print(ast.dump(root_node, indent=4))

注：dump()时设置indent=4（缩进4空格），可以使打印输出的内容更加直观。

输出结果如下：

Module(
    body=[
        Expr(
            value=Call(
                func=Name(id='print', ctx=Load()),
                args=[
                    Constant(value='hello world')],
                keywords=[]))],
    type_ignores=[])

从语法树中可以看出，该语句加载(Load())了名(Name())为print的函数接口(func)，函数传参(args)是值为’hello world’(value)的常量(Constant)。

a = func(1) + func2(func3(3) + func4(1))

第二段代码，显示表达式的抽象语法树：a = func(1) + func2(func3(3) + func4(1))

import ast

deffunc(inputval):
    output = inputval + 1return output

deffunc2(inputval):
    output = inputval + 2return output

deffunc3(inputval):
    output = inputval + 3return output

deffunc4(inputval):
    output = inputval + 4return output

root_node = ast.parse('a = func(1) + func2(func3(3) + func4(1))')
print(ast.dump(root_node, indent=4))

输出如下：

节点分析

通过上述两个例子，可以更好地理解AST的节点构成。节点可以分类为：

常量节点（Literals）

变量节点（Variables）

表达式节点（Expressions）

声明节点（Statements）

控制流节点（Control flow，if/for/while等）

函数和类的定义节点（Function and class definitions）

异步和等待节点（Async and await）

顶层节点（Top level nodes）

Module是AST树的顶层节点，它的body属性以list形式存储了各个节点，同时还有type_ignores属性，记录标志了# type: ignore的所在行。

注：当python指令以exec模式运行时，根节点为Module；以single模式运行，根节点为classInteractive；以eval 模式运行，根节点为Expression。

注：eval() 和 exec() 函数的功能是相似的，都可以执行一个字符串形式的 Python 代码（代码以字符串的形式提供），相当于一个 Python 的解释器。二者不同之处在于，eval() 执行完要返回结果，而 exec() 执行完不返回结果。

Name是一个变量节点，记录变量的名称(id)和调用方式(ctx)。

Assign是赋值声明节点，targets属性中以list存储要被赋值的对象（节点），当存在多个被赋值对象时，每个对象都被赋同一个值。value是单个节点。

BinOp是一个二进制操作的声明节点，需要传入三个参数：left节点，op操作方式和right节点。

Call是一个函数调用的声明节点，需要传入func、args等参数。

其他各个节点的具体介绍，参考文档：Meet the Nodes

使用ast模块修改运行流程

ast模块支持我们在不修改原有代码/模块的情况下，调整代码的执行流程。

比如说，原有模块实现的是一个加法操作，ast模块接收到原有代码的加法操作后，能够自定义修改成减法操作并运行。

参考文档：Working on the Tree

使用ast.NodeVisitor查找节点

要实现抽象语法树的修改，可以使用的工具是ast.NodeVisitor，这是ast里专门用于查找树中节点的工具。

用例如下：

import ast

# 字符串：定义加法函数并执行
FUNC_DEF = \
"""
def add(x, y):
    return x + y
print(add(3, 5))
"""# 解析上面这个字符串，生成抽象语法树
root_node = ast.parse(source=FUNC_DEF)

# 定义一个节点查找类，需要继承ast.NodeVisitor模块classMyNodeVisitor(ast.NodeVisitor):# 查找抽象语法树里的 函数定义 类型的节点defvisit_FunctionDef(self, node):
        print(node.name) # 打印当前节点下的函数名
        self.generic_visit(node) # 遍历子节点defvisit_BinOp(self, node):# 查找抽象语法树里的 二进制操作 类型的节点if isinstance(node.op, ast.Add): # 判断是否出现 加操作
            print('+') # 打印 加操作
        self.generic_visit(node) # 遍历子节点# 实例化 节点查找类，并调用visit接口进行遍历查找
MyNodeVisitor().visit(node=root_node)

输出：

add
+

如代码所示，要查找指定类型的节点，需要执行以下步骤：

1、定义一个查找的类，该类继承ast.NodeVisitor模块

2、定义查找指定类型节点的函数，函数名为visit_xxx(self, node)，xxx为指定类型节点的类型名，如FunctionDef或BinOp，具体有哪些节点参考文档：Meet the Nodes

3、在函数内调用self.generic_visit(node)，从而让函数继续遍历当前节点的子节点

4、实例化我们定义的查找的类，然后调用接口visit(node=xxx)，xxx为抽象语法树的根节点，从而实现遍历查找动作。

使用ast.walk(node)查找节点

或者说，我们也可以使用ast.walk(node)方法遍历所有节点，这个方法类似迭代器操作，但是这个方法不保证遍历顺序是有序的。

import ast

# 字符串：定义加法函数并执行
FUNC_DEF = \
"""
def add(x, y):
    return x + y
print(add(3, 5))
"""# 解析上面这个字符串，生成抽象语法树
root_node = ast.parse(source=FUNC_DEF)

# 使用ast.walk方法遍历root_node里的所有节点for node in ast.walk(node=root_node):
	# 判断是不是FunctionDef节点if isinstance(node, ast.FunctionDef):
        print(f'Find Functiondef:{node.name:s}')
	# 判断是不是BinOp节点if isinstance(node, ast.BinOp):
    	# 判断是不是BinOp节点下的Add方法if isinstance(node.op, ast.Add):
            print('+')

输出：

Find Functiondef:add
+

修改节点里的操作

把BinOp节点中的加法操作改成减法：

import ast
import astunparse

FUNC_DEF = \
"""
def add(x, y):
    return x + y
print(add(3, 5))
"""

root_node = ast.parse(source=FUNC_DEF)

classMyNodeVisitor(ast.NodeVisitor):defvisit_FunctionDef(self, node):
        print(node.name)
        self.generic_visit(node)
    
    defvisit_BinOp(self, node):if isinstance(node.op, ast.Add):
        	# 把加法操作改成减法
            print('+ -> -')
            node.op = ast.Sub()
        self.generic_visit(node)

# 执行抽象语法树的内容
print('\nexec...')
exec(compile(root_node, '', 'exec'))

print('\nvisit...')
MyNodeVisitor().visit(node=root_node)

# 重新执行抽象语法树的内容
print('\nexec...')
exec(compile(root_node, '', 'exec'))

# 把修改后的抽象语法树恢复成代码，打印出来
print(astunparse.unparse(root_node))

输出：

exec...
8

visit...
add
+ -> -

exec...
-2defadd(x, y):return (x - y)
print(add(3, 5))

可以看出，经过visit操作后，加法操作被改成了减法操作，执行该抽象语法树后的加法操作（3+5=8）变成了减法操作（3-5=-2）。

同时，通过unparse方法，还能把修改后的语法树恢复成代码。

替换节点

ast.NodeVisitor方法或ast.walk(node)方法可以对抽象语法树的节点进行遍历，然后在遍历时对节点内部的参数和方法进行修改调整。

但是如果我们想要替换一整个节点，就需要使用另一个方法：ast.NodeTransformer。

这个方法在前者的基础上，还会在visit_xxx函数中返回一个节点变量，返回的这个节点会替换原有的节点。

如果返回的节点是None，那么该位置的节点会被移除。

替换节点需要关注的操作，是节点的创建和插入。

用例如下：

import ast
import astunparse

FUNC_DEF = \
"""
def add(x, y):
    return x + y
print(add(3, 5))
"""

root_node = ast.parse(source=FUNC_DEF)

classMyReplaceNode(ast.NodeTransformer):'''
    修改节点
    '''defvisit_BinOp(self, node):# 寻找二进制操作节点中的加法操作if isinstance(node.op, ast.Add):
            # 新建一个节点，传入参数为原有节点的参数，操作是减法
            my_new_node = ast.BinOp(
                left = node.left,
                op = ast.Sub(),
                right = node.right
            )
            # 新建节点缺少lineno和col_offset属性，使用ast.copy_location接口从旧节点拷贝过来
            my_new_node = ast.copy_location(new_node=my_new_node, old_node=node)
            # 返回新节点return my_new_node
        # 返回旧节点return node

print('\nexec before replace...')
exec(compile(root_node, '', 'exec'))

print('\nerplacce...')

my_replace_node = MyReplaceNode()
my_replace_node.visit(root_node)

print('\nexec after replace...')
exec(compile(root_node, '', 'exec'))

print('\nunparse code...')
print(astunparse.unparse(root_node))

输出：

exec before replace...
8

erplacce...

exec after replace...
-2

unparse code...


defadd(x, y):return (x - y)
print(add(3, 5))

新建节点时，节点的lineno和col_offset这两个属性需要关注下，我们手动创建的节点默认不带这两个属性，但是ast解析的语法树中的节点携带，且需要这两个属性。我们有以下三种方法配置这两个属性：

1、ast.fix_missing_locations(node) ：从父节点node复制这两个属性的值，然后递归地查找子节点中缺少这两个属性的位置，填充父节点的值。这是一种粗暴但是直接的方法。

2、ast.copy_location(new_node, old_node)：从old_node节点拷贝这两个属性的值，填充至new_node节点中，然后返回new_node。做节点替换操作时这个操作会很好用。

3、ast.increment_lineno(node, n=1)：将节点node及其子节点从起始行号到结束行号递增n。当需要将代码“移动”到文件中的不同位置时，这个操作非常有用。