修改多个节点 python ast.NodeTransformer

Question

我有这样的输入源码

def foo(my_input):
  return my_input + 42

想要这样变形

def method_name(arg0):
  return my_input + 42

为此目的的ast节点转换器是这样写的。

class MyRenamer(ast.NodeTransformer):

  def __init__(self):
    self._arg_count = 0

  def visit_FunctionDef(self, node):
    node.name = "method_name"
    return node

  def visit_arg(self, node):
    node.arg = "arg_{}".format(self._arg_count)
    self._arg_count += 1
    return node

但是当我这样调用上面的转换器时

node = ast.parse(code)
renamer = MyRenamer()
node2 = renamer.visit(node)
print(astor.to_source(node2))

我得到的输出是

def method_name(my_input):
    return my_input + 42

这里函数的参数没有改变。

Answer 1

访问者需要通过访问当前访问节点的所有children来遍历AST。方法 generic_visit() 会为您完成此操作，但您必须在每个 visit_... 方法中调用它，或者至少对于可能 children 的方法。

import ast
import astor

class MyRenamer(ast.NodeTransformer):

  def __init__(self):
    self._arg_count = 0

  def visit_FunctionDef(self, node):
    node.name = "method_name"
    self.generic_visit(node)
    return node

  def visit_arg(self, node):
    node.arg = "arg_{}".format(self._arg_count)
    self._arg_count += 1
    self.generic_visit(node)
    return node

code = """
def foo(my_input):
  return my_input + 42
"""

node = ast.parse(code)
renamer = MyRenamer()
node2 = renamer.visit(node)
print(astor.to_source(node2))

def method_name(arg_0):
    return my_input + 42

这给出了您预期的输出，但在更大的上下文中会将所有函数重命名为“method_name”，这可能不是您想要的。而且函数体中还有标识符，大概也需要重命名。

修改多个节点 python ast.NodeTransformer

Modify multiple nodes with python ast.NodeTransformer

python

abstract-syntax-tree