Python AST NodeTransformer:Return 多个节点

Python AST NodeTransformer: Return multiple nodes

在使用ast.NodeTransformer时,有没有一种简单的方法可以return多个节点来替换单个节点?例如,假设我想重写

形式的所有表达式

f(g())_x1 = g(); g(_x1)

如果 visit_Expr 可以 return 多个而不是单个节点,那么这样做会很容易。不过我似乎无法让它工作,所以我认为这是不是的方法。任何建议将不胜感激。

[Update] 作为更新,我有一个工作版本,它在列表中累积新旧节点,并将它们分配给封闭范围节点的主体(例如 ForWhileModule 节点等)。这绝对是一种 hacky 的方式,并且怀疑有更好的方法。我会保留它,以防有人知道。

[final update] 查看 NodeTransformer 的文档实际上完全有可能 return 节点列表,如果节点是语句集合的一部分。

对于声明节点,您可以return一个新节点列表。这使您可以用多个语句替换单个语句。引用文档:

For nodes that were part of a collection of statements (that applies to all statement nodes), the visitor may also return a list of nodes rather than just a single node.

对于您的表达式,有一个顶级 Expr() 表达式语句节点:

>>> ast.dump(ast.parse('f(g())'))
"Module(body=[Expr(value=Call(func=Name(id='f', ctx=Load()), args=[Call(func=Name(id='g', ctx=Load()), args=[], keywords=[])], keywords=[]))])"
>>> import ast
>>> ast.dump(ast.parse('f(g())'))
"Module(body=[Expr(value=Call(func=Name(id='f', ctx=Load()), args=[Call(func=Name(id='g', ctx=Load()), args=[], keywords=[])], keywords=[]))])"

因此,您所要做的就是提供一个 visit_Expr 处理程序,该处理程序 return 是一个包含 2 个语句节点的列表;第一个调用 g() 的赋值(Assign() 语句节点),第二个调用 f 传入新变量名(另一个 Expr() 语句节点) .

我会在您进入和离开表达式语句上下文时设置标志的转换器子类中保持状态,并跟踪该上下文下的调用堆栈。然后当您 return 到 visit_Expr 时,您 return 您的新设置:

self._expr_statement = False

def visit_Expr(node):
    self._expr_statement = True
    self.generic_visit(node)
    self._expr_statement = False

    if <specific state on self matches expectations>:
        tempvar = <generated_new_name>
        return [
            Assign([Name(tempvar, Store())], <inner_call>),
            Expr(Call(
                <outer_function_name_expr>,
                args=[Name(tempvar, Load())],
                keyword=[]))
        ]
    else:
        # no replacement takes place
        return node

NodeTransformer then uses that list of elements to replace the previousExpr()` 节点。

请注意,您仍然必须调用 self.generic_visit(node) 以确保仍处理嵌套节点;然后将为这些嵌套节点调用进一步的 visit_* 方法。这样就可以在 visit_Call 方法中检查 self._expr_statement 标志,然后测试是否存在嵌套调用,然后在 self 上为 [=25= 存储足够的上下文] 方法 return.