ast 模块的 documentation说明如何使用 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 或使用 ast.unparse适用于 Python 3.9 及更高版本。
Python3 解析脚本** 可以在 Tools 中找到CPython 存储库的目录。在 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)])])"
*归功于this answer用于记录未解析脚本的存在。
** 使用与正在使用的 Python 版本相对应的 git 分支中的脚本版本。例如,由于版本各自的语法差异,在 3.7 代码上使用 3.6 分支的脚本可能会失败。
关于python - 向抽象语法树中插入一个节点,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/46388130/