【python】ast模块介绍和使用

ast模块

官方文档：ast — Abstract Syntax Trees
教程文档：Getting to and from ASTs

ast模块简介

参考文章：python compiler.ast_Python Ast介绍及应用
Python官方提供的CPython解释器对python源码的处理过程如下：

1. Parse source code into a parse tree (Parser/pgen.c)
2. Transform parse tree into an Abstract Syntax Tree (Python/ast.c)
3. Transform AST into a Control Flow Graph (Python/compile.c)
4. Emit bytecode based on the Control Flow Graph (Python/compile.c)

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

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

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

查看ast抽象语法树

print(‘hello world’)

查看print(‘hello world’)语句的抽象语法树：

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

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

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

<ast.Module object at 0x0000021443614940>
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

def func(inputval):
    output = inputval + 1
    return output

def func2(inputval):
    output = inputval + 2
    return output

def func3(inputval):
    output = inputval + 3
    return output

def func4(inputval):
    output = inputval + 4
    return 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模块
class MyNodeVisitor(ast.NodeVisitor):
    # 查找抽象语法树里的 函数定义 类型的节点
    def visit_FunctionDef(self, node):
        print(node.name) # 打印当前节点下的函数名
        self.generic_visit(node) # 遍历子节点
    
    def visit_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)

class MyNodeVisitor(ast.NodeVisitor):
    def visit_FunctionDef(self, node):
        print(node.name)
        self.generic_visit(node)
    
    def visit_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...
-2

def add(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)

class MyReplaceNode(ast.NodeTransformer):
    '''
    修改节点
    '''
    def visit_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...


def add(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。当需要将代码“移动”到文件中的不同位置时，这个操作非常有用。