python ast模块使用

ast(Abstract Syntax Trees)是python中非常有用的一个模块，我们可以通过分析python的抽象语法树来对python的代码进行分析和修改。
ast作用在python代码的语法被解析后，被编译成字节码之前。

一. ast

1. 获取语法树

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

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

通过ast的parse方法得到ast tree的根节点root_node, 我看可以通过根节点来遍历语法树，从而对python代码进行分析和修改。
ast.parse(可以直接查看ast模块的源代码)方法实际上是调用内置函数compile进行编译，如下所示：

def parse(source, filename='<unknown>', 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)

传递给compile特殊的flag = PyCF_ONLY_AST, 来通过compile返回抽象语法树。
（关于compile的第四个参数flag有时间可以深入研究下）

2. 节点类型分析

语法树中的每个节点都对应ast下的一种类型，根节点是ast.Moudle类型，在分析的时候可以通过isinstance函数方便的进行节点类型的判断。
ast中存在的节点的所有类型可以参考：ast节点类型
比如 a = 10这样一条语句对应ast.Assign节点类型，而Assign节点类型分别有两个子节点， 分别为ast.Name类型的a和ast.Num类型的10等。
我们可以通过ast.dump(node)函数来将node格式化，并进行打印，以查看节点内容，以“a = 10”这行代码为例。

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

(1) root节点

Module(body=[Assign(targets=[Name(id='a', ctx=Store())], value=Num(n=10))])
root节点是Module类型，由于只有一行代码，所有root节点只有Assign这样一个子节点。

(2) 子节点

Assign(targets=[Name(id='a', ctx=Store())], value=Num(n=10))
上述的Assign节点有三个子节点，分别是Name, Store和Num.
Name(id='a', ctx=Store())
Num(n=10)
而Name有一个子节点，Store.
Store()(Store表示Name中操作时赋值, 类型的有Load，del, 具体参考节点类型的文档)
一个简单的“a = 10”的这样一行代码，我们就可以通过上述的这种ast tree去分析和修改代码结构。

二. 语法树的遍历分析

可以通过ast模块的提供的visitor来对语法树进行遍历。
ast.NodeVisitor是一个专门用来遍历语法树的工具，我们可以通过继承这个类来完成对语法树的遍历以及遍历过程中的处理。

1. visitor的定义

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类型"。
visitor首先从根节点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)这段代码，否则visitor不会继续访问当前节点的子节点。
e.g. 如果定义如下的函数：
def visit_Moudle(self, node):
     print type(node).__name__
那么，首先访问根节点root，root为Moudle类型，会调用visit_Moudle函数，由于visit_Moudle函数中没有调用NodeVisitor.generic_visit(self, node)，所以此次遍历只遍历了根节点root，并没有遍历其他节点。

2. walk方式遍历

也可以通过ast.walk对ast tree进行遍历，如下：

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

三. 节点的修改

ast模块同样提供了一个NodeTransfomer节点来支持对node的修改，NodeTransfomer继承自NodeVisitor，并重写了generic_visit函数。
对于NodeTransfomer的generic_visit以及visit_ + 节点类型的函数，都需要返回一个node，可以返回原始node，一个新的替代的node，或者是返回Node代表remove掉这个节点。
假设我们有如下的代码：

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

我们定义一个NodeTransform的visitor如下：

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以及print a的时候，都会将a替换成a_rep，并返回一个新节点。
在visit_Num中，简单粗暴的将10替换成了100，返回修改后的原节点。
我们通过如下方式运用这个NodeTransfomer visitor：

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, "<string>", "exec")
exec code_object

ast作用在python解析语法之后，编译成pyCodeObject字节码结构之前，通过NodeTransformer修改后，返回修改后的语法树，我们通过内置模块compile编译成pyCodeObject对象，交给python虚拟机执行。
执行结果：100
可以看到，我们同时将a = 10和print a两处将a名字换成了a_rep，并将10替换成了100，最后打印的结果是100，成功修改了语法树的节点。
关于节点的修改，这里有比较好的例子可以参考：https://greentreesnakes.readthedocs.org/en/latest/examples.html
注意：
修改语法树节点，尤其是删除一个语法树节点时要慎重，因为修改或者删除后有可能返回错误的语法树，直到compile或者执行的时候才会发现问题。
通过节点修改python code就可以通过上述方法进行，不过请注意，在运用visitor的代码中有ast.fix_missing_locations(root)这样一行代码，这是因为我们自己创建的节点是不包含lineno以及col_offset这些必要的属性，必须手动修改添加指定，新添加的节点代码的行位置以及偏移位置。

四. 修复节点位置

 

1  属性分析

每个节点都有一些相应的属性，lineno以及col_offset是每个节点都必须有的属性，分别代表行号以及在这行中的偏移。
另外每个节点都有一些自己的特殊属性，如上诉的Module含有body属性，Assign含有targets属性等。
lineno以及col_offset这两个属性，如果是python中原本代码的节点，如Assign、Name、Num等（注：Moudle和Store这样的节点是没有lineno以及col_offset属性的），但是如果我们通过NodeTransFormer新增的节点，默认是不存在这些属性的，我们可以通过三种方法来fix这些节点的lineno以及col_offset属性。

2  属性的修复

我们可以通过相应的方法，对默认没有lineno以及col_offset的节点进行位置的修复，以方便在代码中获取每个节点的位置信息，主要有三种方法进行修复。
1）ast.fix_missing_locations(node)
函数递归的将父节点的位置信息（lineno以及col_offset）赋值给没有位置信息的子节点。
2）ast.copy_location(new_node, node)
将node的位置信息拷贝给new_node节点，并返回new_node节点。当我们将旧节点替换成一个新节点的时候，这种方法比较适用。
3）ast.increment_lineno(node, n=1)
将node节点以及其所以子节点的行号加上n。

3  分析

我们通过“三. 节点的修改"中的例子来分析location信息。
在例子中，我们只有在visit_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将旧节点的location信息拷贝给新节点。

参考：
Green Tree Snakes - the missing Python AST docs --非常详尽的ast的模块的分析文档。
Simple example of how to use ast.NodeVisitor --stackoverflow上一篇比较明了的回答
Instrumenting the AST --简单的分析与应用
ast — Abstract Syntax Trees --官方文档
AST 模块：用 Python 修改 Python 代码 --比较详细的介绍了修改

 

转载于https://blog.youkuaiyun.com/ma89481508/article/details/56017697