AST 模块：用 Python 修改 Python 代码

本文介绍如何利用Python的AST模块得到或者修改python虚拟机编译过程中生成的语法树

CPython 的编译过程

在开始之前，我们应该先看看 CPython 的编译过程，这个过程在 PEP 339 中有详细的描述。

当然，在读这篇文章的时候，你并不需要对这个步骤有很深入的理解，不过这可以帮助你对整个过程有一个大体的了解。

首先，编译器会根据源代码生成一棵语法分析树 (Parse Tree)，随后，再根据语法分析树建立抽象语法树 (AST, Abstract Syntax Tree)。从 AST 中可以生成出控制流图 (CFG, Control Flow Graph)，最后再将控制流图编译为代码对象 (Code Object)。

图中标蓝的部分就是 AST 这一步，也就是我们今天所关注的部分。Python 从 2.6 开始就提供了现在这样的 ast 模块，它提供了一种访问和修改 AST 的简单方式。

通过这个，我们可以从 AST 中生成代码对象，也可以出于某些原因，根据修改过的 AST 重新生成源代码。

创建 AST

先来写一点简单的代码，我们写一个叫做 add 的函数，然后观察它所生成的 AST。

>>> import ast
>>> expr = """
... def add(arg1, arg2):
...     return arg1 + arg2
... """
>>> expr_ast = ast.parse(expr)
>>> expr_ast
<_ast.Module object at 0x10a7a09d0>

现在我们已经生成了一个 ast.Module 对象，我们来看看它的内容：

>>> ast.dump(expr_ast)
"Module(
    body=[
        FunctionDef(
            name='add', args=arguments(
                args=[
                    Name(id='arg1', ctx=Param()),
                    Name(id='arg2', ctx=Param())
                    ],
                    vararg=None,
                    kwarg=None,
                    defaults=[]),
            body=[
                Return(
                    value=BinOp(
                        left=Name(id='arg1', ctx=Load()),
                        op=Add(),
                        right=Name(id='arg2', ctx=Load())))
                ],
            decorator_list=[])
    ])"

正如我们所见， Module 是父节点，它的 body 中包含了一个函数定义的元素，这个函数定义包含了函数名、参数列表和函数体。函数体又包含了一个单独的 Return 节点，节点中含有一个 Add 运算。

修改 AST

我们如何修改这棵树以改变代码的作用呢？为了说明这个问题，我们来做点也许你永远也不会在你自己代码中做的疯狂的事情吧。我们将遍历这棵树，并且将 Add 运算修改为 Mult 运算。看，我说过这很疯狂吧！

我们要先建立一个 NodeTransformer 变换器的子类，并且定义 visit_BinOp 方法。每当这个变换器访问到一个二元运算符节点时，就会调用这个方法。

class CrazyTransformer(ast.NodeTransformer):

    def visit_BinOp(self, node):
        print node.__dict__
        node.op = ast.Mult()
        print node.__dict__
        return node

现在我们已经定义好了我们这个奇怪的变换器，让我们看看将它应用于我们开始时写的那些代码会怎么样：

>>> transformer = CrazyTransformer()
>>> transformer.visit(expr_ast)
{
    'op': <_ast.Add object at 0x10a8321d0>,
    'right': <_ast.Name object at 0x10a839390>,
    'lineno': 3, 'col_offset': 8,
    'left': <_ast.Name object at 0x10a839350>}
{
    'op': <_ast.Mult object at 0x10a839510>,
    'right': <_ast.Name object at 0x10a839390>,
    'lineno': 3, 'col_offset': 8,
    'left': <_ast.Name object at 0x10a839350>}

你可以从输出的结果对比发现， Add 节点已经被替换成了一个 Mult 。我们有许多方法没有提到，比如访问子节点，不过这个例子已经足以刻画出它的基本原理。

编译和执行修改后的 AST

我们在最初的代码后面加上一个调用，比如：

print add(4, 5)

让我们看看这些代码是如何运行的：

>>> unmodified = ast.parse(expr)
>>> exec compile(unmodified, '<string>', 'exec')
9
>>> transformer = CrazyTransformer()
>>> modified = transformer.visit(unmodified)
>>> exec compile(modified, '<string>', 'exec')
20

我们可以看到，未修改的和修改后的 AST 所编译出的代码，一个输出了9，一个输出了20。

重新翻译回源代码

最后，我们可以用 unparse 模块将修改后的代码转换回对应的源代码， unparse 模块可以在 这里 找到。

>>> unparse.Unparser(modified, sys.stdout)

def add(arg1, arg2):
    return (arg1 * arg2)
print add(4, 5)

正如我们所看到的， * 运算符取代了 + 。在这个反解析工具对于理解你的 AST 变换器如何修改代码很有帮助。