深入解析Python中的AST模块：代码分析与修改的强大工具

ast（抽象语法树）是Python中非常有用的模块。我们可以通过分析python的抽象语法树来分析和修改python代码。

ast 在 python 代码语法解析之后、编译为字节码之前使用。

1. ast 1. 获取语法树

ast模块的基本使用非常简单。我们可以通过以下代码快速得到一棵抽象语法树：

import ast
root_node = ast.parse("print 'hello world'")
root_node -> <_ast.Module object at 0x9e3df6c>

通过ast的parse方法得到ast树的根节点root_node。我觉得可以通过根节点遍历语法树来分析和修改python代码。

ast.parse（可以直接查看ast模块的源码）方法实际上调用了内置函数compile进行编译，如下所示：

def parse(source, filename='', mode='exec'):
    """
    Parse the source into an AST node.
    Equivalent to compile(source, filename, mode, PyCF_ONLY_AST).
    """
    return compile(source, filename, mode, PyCF_ONLY_AST)

传递特殊标志 = PyCF_ONLY_AST 进行编译，通过编译返回抽象语法树。

（有时间可以深入研究compile的第四个参数flag）

2. 节点类型分析

语法树中的每个节点对应ast下的一个类型。根节点的类型为ast.Moudle。分析时，通过isinstance函数可以轻松确定节点类型。

ast中存在的所有节点类型可以参考：ast节点类型

例如a=10这样的语句对应的是ast.Assign节点类型，Assign节点类型有两个子节点，分别是ast.Name类型的a和ast.Num类型的10。

我们可以使用ast.dump(node)函数格式化节点并打印来查看节点内容。以代码行“a = 10”为例。

Module(body=[Assign(targets=[Name(id='a', ctx=Store())], value=Num(n=10))])

(1)根节点

模块(body=[分配(目标=[名称(id='a', ctx=Store())], 值=Num(n=10))])

根节点是Module类型。由于只有一行代码，所以所有根节点只有一个子节点Assign。

(2)子节点

分配（目标= [名称（id ='a'，ctx = Store（））]，值= Num（n = 10））

上面的Assign节点有3个子节点，分别是Name、Store和Num。

名称(id='a', ctx=Store())

数量(n=10)

Name 有一个子节点 Store。

Store()（Store表示Name中操作时赋值的值，类型有Load和del，具体请参考节点类型的文档）

一行简单的代码如“a = 10”，我们就可以通过上面的ast树来分析并修改代码结构。

2. 语法树遍历分析

可以通过ast模块提供的visitor来遍历语法树。

ast.NodeVisitor是一个专门用来遍历语法树的工具。我们可以继承这个类来完成语法树的遍历以及遍历过程中的处理。

1. 访客的定义

class CodeVisitor(ast.NodeVisitor):
    def generic_visit(self, node):
        print type(node).__name__
        ast.NodeVisitor.generic_visit(self, node)
 
    def visit_FunctionDef(self, node):
        print type(node).__name__
        ast.NodeVisitor.generic_visit(self, node)
 
    def visit_Assign(self, node):
        print type(node).__name__

如上面的代码所示，定义了类CodeVisitor，继承自NodeVisitor。函数主要有两种类型，一种是generic_visit，另一种是“visit_”+“Node type”。

访问者首先从根节点root开始遍历。遍历过程中，假设节点类型为Assign。如果存在visit_Assign类型的函数，则调用visit_Assgin函数。如果不存在，则调用 generic_visit 函数。

一般来说，每种节点类型都有专用的类型处理函数。如果不存在，则调用通用处理函数generic_visit。

关于visitor的语法树遍历，stackoverflow上有一篇文章讲得比较详细：Simple example of how to use ast.NodeVisitor

注意：

在每个函数处理中，根据需要添加代码ast.NodeVisitor.generic_visit(self, node)，否则访问者将不会继续访问当前节点的子节点。

例如，如果您定义以下函数：

def Visit_Moudle(自身, 节点):

打印类型（节点）.__name__

然后，首先访问根节点root，它是Moudle类型，会调用visit_Moudle函数。由于visit_Moudle函数中没有调用NodeVisitor.generic_visit(self, node)，因此本次遍历只遍历了根节点root，并没有遍历其他节点。

2.步行模式遍历

还可以通过ast.walk遍历ast树，如下：

for node in ast.walk(tree):
    if isinstance(node, ast.FunctionDef):
        print(node.name)

3、节点修改

ast模块还提供了NodeTransfomer节点来支持节点的修改。 NodeTransfomer继承自NodeVisitor，并重写了generic_visit函数。

对于NodeTransfomer的generic_visit和visit_+节点类型函数，都需要返回一个节点。您可以返回原始节点、新的替换节点或返回 Node 来表示删除该节点。

假设我们有以下代码：

"""ast test code"""
a = 10
b = "test"
print a

我们定义一个 NodeTransform 访问者如下：

class ReWriteName(ast.NodeTransformer):
    def generic_visit(self, node):
        has_lineno = getattr(node, "lineno", "None")
        col_offset = getattr(node, "col_offset", "None")
        print type(node).__name__, has_lineno, col_offset
        ast.NodeTransformer.generic_visit(self, node)
        return node
 
    def visit_Name(self, node):
        new_node = node
        if node.id == "a":
            new_node = ast.Name(id = "a_rep", ctx = node.ctx)
        return new_node
 
    def visit_Num(self, node):
        if node.n == 10:
            node.n = 100
        return node

在visit_Name 中，将变量“a”替换为变量“a_rep”。当执行a = 10并打印a时，a将被a_rep替换并返回一个新节点。

在visit_Num中，只需将10替换为100并返回修改后的原始节点。

我们通过以下方式使用这个 NodeTransfomer 访问者：

file = open("code.py", "r")
source = file.read()
visitor = ReWriteName()
root = ast.parse(source)
root = visitor.visit(root)
ast.fix_missing_locations(root)
 
code_object = compile(root, "", "exec")
exec code_object

ast 在 python 解析语法之后、将其编译成 pyCodeObject 字节码结构之前使用。经过NodeTransformer修改后，返回修改后的语法树。我们通过内置模块compile将其编译成pyCodeObject对象，交给python虚拟机执行。

执行结果：100

可以看到，我们把a=10同时打印a，将a名称改为a_rep，并将10替换为100，最终打印结果为100，语法树的节点修改成功。

关于节点的修改，这里有一些很好的例子可供参考：

注意：

修改语法树节点时要小心，尤其是删除语法树节点时，因为修改或删除后可能会返回错误的语法树，直到编译或执行时才会发现问题。

通过节点修改python代码可以通过上面的方法完成，但是请注意，使用visitor的代码中有一行代码如ast.fix_missing_locations(root)。这是因为我们创建的节点不包含lineno和col_offset。必须手动修改属性，添加新添加节点代码的指定、行位置和偏移位置。

4.修复节点位置

1 属性分析

每个节点都有一些相应的属性。 lineno 和 col_offset 是每个节点必须具有的属性，分别表示行号和在该行中的偏移量。

另外，每个节点都有一些自己特殊的属性，比如调用的Module包含body属性，Assign包含targets属性等。

lineno和col_offset这两个属性，如果是python中原代码的节点，如Assign、Name、Num等（注：Moudle、Store等节点没有lineno和col_offset属性），但是如果我们添加它们通过 NodeTransFormer 节点，这些属性默认不存在。我们可以通过三种方法修复这些节点的lineno和col_offset属性。

2.属性修复

我们可以使用相应的方法来修复默认没有lineno和col_offset的节点的位置，以方便在代码中获取每个节点的位置信息。修复方法主要有三种。

1）ast.fix_missing_locations（节点）

该函数递归地将父节点的位置信息（lineno和col_offset）分配给没有位置信息的子节点。

2) ast.copy_location(new_node, 节点)

将节点的位置信息复制到new_node节点并返回new_node节点。当我们用新节点替换旧节点时，这种方法更适合。

3）ast.increment_lineno（节点，n=1）

该节点node及其所有子节点的行号加n。

3 分析

我们通过“3.节点修改”​​中的例子来分析位置信息。

在示例中，我们仅在访问期间返回新节点_Name。此时，该节点不具有lineno和col_offset属性。我们可以通过两种方式获得它们。

首先如上面的代码，使用ast.fix_missing_locations函数来修复。在“a = 10”和“print a”中，Name节点a与父节点的lineno相同，但此时col_offset会不同。

其次，我们修改visit_Name的代码如下：

def visit_Name(self, node):
    new_node = node
    if node.id == "a":
        new_node = ast.Name(id = "a_rep", ctx = node.ctx)
        ast.copy_location(new_node, node)
    return new_node

通过copy_location将旧节点的位置信息复制到新节点。

参考：

Green Tree Snakes - 缺失的 Python AST 文档 - 非常详细的 ast 模块分析文档。

如何使用ast.NodeVisitor的简单示例——stackoverflow上一篇文章中比较清晰的答案

检测 AST——简单分析和应用

ast — 抽象语法树--官方文档

--更详细地介绍了修改内容