CODEWITHSUNDEEP
pythonabstract-syntax-tree
Jayendra Parmar
Jayendra Parmar

Reputation: 774

Modify multiple nodes with python ast.NodeTransformer

I have a input source code like this

def foo(my_input):
  return my_input + 42

and want it to transform like this

def method_name(arg0):
  return my_input + 42

The ast node transformer for this purpose is written like this.

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

but when I call the above transformer like this.

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

The output I get is

def method_name(my_input):
    return my_input + 42

Here the argument of the function has not been changed.

Upvotes: 2

Views: 1735

Answers (1)

cbare
cbare

Reputation: 12468

The visitor needs to walk the AST by visiting all children of the currently visited node. The method generic_visit() does that for you, but you have to call it in every visit_... method, or at least for those where children are a possibility.

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

This gives your expected output, but in a larger context will rename all functions to "method_name", which might not be desired. And there's still the identifier in the function body that presumably also needs to be renamed.

Upvotes: 2

Related Questions