将节点插入抽象语法树 [英] Insert a node into an abstract syntax tree

问题描述

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

例如，给定这个模块:

import foo进口酒吧类巴兹(对象):定义垃圾邮件(自我):经过

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

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

解决方案

Python AST 本质上是由嵌套列表组成的，因此新节点可以在构建后插入到这些列表中.

首先获取要改变的AST:

<预><代码>>>>root = ast.parse(open('test.py').read())>>>ast.dump(根)"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 属性，它包含了模块的顶级元素:

<预><代码>>>>根体[<_ast.Import 对象在 0x7f81685385d0>，<_ast.Import 对象在 0x7f8168538950>，<_ast.ClassDef 对象在 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 对象在 0x7f816812ef90>

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

Python3 解析脚本^** 可以在 Tools 目录.在 Python2 中，它位于 Demo 目录中.

<预><代码>>>>from unparse import Unparser>>>buf = StringIO()>>>解析器(根，buf)<unparse.Unparser 实例在 0x7f81685c6248>>>>buf.seek(0)>>>打印(buf.read())进口富进口酒吧进口quux类巴兹(对象):鸡蛋 = '火腿'定义垃圾邮件(自我):经过>>>

使用ast.unparse:

<预><代码>>>>未解析 = ast.unparse(root)>>>打印(未解析)

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

<预><代码>>>>root = ast.parse('import foo')>>>ast.dump(根)"模块(body=[Import(names=[alias(name='foo', asname=None)])])";

^* 归功于 这个答案 记录了未解析脚本的存在.>

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

The ast module's documentation explains how to replace a node in the AST using the NodeTransformer class, but does not explain how to insert a new node into the tree.

For example, given this module:

import foo
import bar

class Baz(object):

    def spam(self):
        pass

I would like to add another import, and set a class variable on Baz.

How can I create and insert these nodes into the AST?

解决方案

Python ASTs are essentially composed of nested lists, so new nodes can be inserted into these lists once they have been constructed.

First, get the AST that is to be changed:

>>> 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=[])])"

We can see that the outer Module has a body attribute that contains the top level elements of the module:

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

Construct an import node and insert:

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

Like the root module node, the class definition node has a body attribute that contains its members:

>>> 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=[])"

So we construct an assignment node and insert it:

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

To finish, fix up line numbers:

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

We can verify that our nodes are in place by dumping the root node with ast.dump, using the unparse^* tool from the CPython repository to generate source from the AST or using ast.unparse for Python 3.9 and later.

The Python3 unparse script^** can be found in the Tools directory of the CPython repository. In Python2 it was located in the Demo directory.

>>> 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
>>> 

Using ast.unparse:

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

When constructing AST nodes, you can get an idea of what the node should look like by using ast.parse and ast.dump (observe that ast.parse wraps the statement in a module):

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

^* Credit to this answer for documenting the existence of the unparse script.

^** Use the version of the script from the git branch that corresponds to the Python version being used. For example, using the script from the 3.6 branch on 3.7 code may fail due to differences in the versions' respective grammars.

这篇关于将节点插入抽象语法树的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！