ast模块的文档介绍了如何使用NodeTransformer类替换AST中的节点,但没有说明如何将新节点插入树中。
例如,给定以下模块:
import foo
import bar
class Baz(object):
def spam(self):
pass
我想添加另一个导入,并在Baz
上设置一个类变量。
我该如何创建并将这些节点插入AST中?
ast模块的文档介绍了如何使用NodeTransformer类替换AST中的节点,但没有说明如何将新节点插入树中。
例如,给定以下模块:
import foo
import bar
class Baz(object):
def spam(self):
pass
我想添加另一个导入,并在Baz
上设置一个类变量。
我该如何创建并将这些节点插入AST中?
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=[])])"
我们可以看到外部模块有一个名为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生成源代码,或者对于Python 3.9及更高版本,使用ast.unparse。
Python3的unparse脚本**可以在CPython存储库的Tools目录中找到。 在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.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)])])"
* 感谢这个回答记录了unparse脚本的存在。
** 使用与正在使用的Python版本相对应的git分支中的脚本版本。例如,在3.7代码中使用来自3.6分支的脚本可能会因版本的语法差异而失败。