向抽象语法树中插入一个节点

Question

ast模块的documentation explains how to replace a node in the AST using the NodeTransformerclass，但没有说明如何向树中插入新节点。

例如，给定此模块：

import foo
import bar

class Baz(object):

    def spam(self):
        pass

我想添加另一个导入，并在 Baz 上设置一个 class 变量。

如何创建这些节点并将其插入到 AST 中？

Answer 1

Python AST 本质上是由嵌套列表组成的，因此一旦构建了新节点就可以将它们插入到这些列表中。

首先，获取要更改的AST：

>>> root = ast.parse(open('test.py').read())

>>> ast.dump(root)
"Module(body=[Import(names=[alias(name='foo', asname=None)]), Import(names=[alias(name='bar', asname=None)]), ClassDef(name='Baz', bases=[Name(id='object', ctx=Load())], body=[FunctionDef(name='spam', args=arguments(args=[Name(id='self', ctx=Param())], vararg=None, kwarg=None, defaults=[]), body=[Pass()], decorator_list=[])], decorator_list=[])])"

我们可以看到外层Module有一个body属性包含了模块的顶层元素：

>>> root.body
[<_ast.Import object at 0x7f81685385d0>, <_ast.Import object at 0x7f8168538950>, <_ast.ClassDef object at 0x7f8168538b10>]

构造导入节点并插入：

>>> import_node = ast.Import(names=[ast.alias(name='quux', asname=None)])
>>> root.body.insert(2, import_node)

与根模块节点一样，class 定义节点有一个包含其成员的 body 属性：

>>> classdef = root.body[-1]
>>> ast.dump(classdef)
"ClassDef(name='Baz', bases=[Name(id='object', ctx=Load())], body=[FunctionDef(name='spam', args=arguments(args=[Name(id='self', ctx=Param())], vararg=None, kwarg=None, defaults=[]), body=[Pass()], decorator_list=[])], decorator_list=[])"

所以我们构造一个赋值节点并插入：

>>> assign_node = ast.Assign(targets=[ast.Name(id='eggs', ctx=ast.Store())], value=ast.Str(s='ham')) 
>>> classdef.body.insert(0, assign_node)

最后，修正行号：

>>> ast.fix_missing_locations(root)
<_ast.Module object at 0x7f816812ef90>

我们可以使用 CPython 存储库中的 unparse^* 工具通过 ast.dump 转储根节点来验证我们的节点是否就位从 AST 或使用 Python 3.9 及更高版本的 ast.unparse 生成源代码。

Python3 解析脚本^** 可以在 Tools directory of the CPython repository. In Python2 it was located in the Demo 目录中找到。

>>> from unparse import Unparser
>>> buf = StringIO()
>>> Unparser(root, buf)
<unparse.Unparser instance at 0x7f81685c6248>
>>> buf.seek(0)
>>> print(buf.read())

import foo
import bar
import quux

class Baz(object):
    eggs = 'ham'

    def spam(self):
        pass
>>>

使用ast.unparse:

>>> unparsed = ast.unparse(root)
>>> print(unparsed)

构造AST节点时，可以通过使用ast.parse和ast.dump来了解节点应该是什么样子（注意ast.parse将语句包装在模块中） :

>>> root = ast.parse('import foo')
>>> ast.dump(root)
"Module(body=[Import(names=[alias(name='foo', asname=None)])])"

^* 感谢 this answer 记录了反解析脚本的存在。

^** 使用与正在使用的 Python 版本相对应的 git 分支中的脚本版本。例如，在 3.7 代码上使用 3.6 分支的脚本可能会由于版本各自语法的差异而失败。

向抽象语法树中插入一个节点

Insert a node into an abstract syntax tree

python

abstract-syntax-tree

python-2.7

python-3.x