将节点插入抽象语法树

Question

ast模块的文档解释了如何使用NodeTransformer类替换AST中的节点,但没有解释如何将新节点插入树中.

例如,给定此模块:

import foo
import bar

class Baz(object):

    def spam(self):
        pass

我想添加另一个导入,并设置一个类变量Baz.

如何创建这些节点并将其插入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)

与根模块节点一样,类定义节点具有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>

我们可以通过转储根节点来验证我们的节点是否到位ast.dump,或者使用CPython存储库中的unparse ^*工具从AST生成源.

可以在CPython存储库的Tools目录中找到Python3 unparse脚本^**.在Python2中,它位于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节点时,您可以通过使用来了解节点应该是什么样的(ast.parse并ast.dump观察它ast.parse在模块中包装语句):

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

^*感谢这个答案用于记录的unparse脚本的存在.

^**使用git分支中与正在使用的Python版本对应的脚本版本.例如,由于版本各自的语法不同,使用3.7代码中3.6分支的脚本可能会失败.