What is ast.NodeTransformer in Python?

ast.NodeTransformer is a subclass of ast.NodeVisitor that walks through an ast tree and can apply transformations to any node in the tree.

NodeTransformer allows you to override the visit methods, and modifying the method’s return value allows you to transform the node as the programmer wishes. If the return value is set to None, then that node is removed from the tree. If no replacement is done to the return value, the return value will be the original node.

Abstract Syntax Tree

Example

import ast
from pprint import pprint
class TransformName(ast.NodeTransformer):
    def visit_Name(self, node):
        return ast.Subscript(
            value=ast.Name(id='data', ctx=ast.Load()),
            slice=ast.Constant(value=node.id),
            ctx=node.ctx
        )
tree = ast.parse('foo', mode='eval')
print("Before transformation:")
pprint(ast.dump(tree))
print("After transformation:")
ast.fix_missing_locations(TransformName().visit(tree))
pprint(ast.dump(tree))

Explanation

The above program modifies any Name node, like variable names, to data[<variable_name>]. Specifically, the above code transforms the node lookup of foo to data['foo'].

  • We define a TransformName class that extends the ast.NodeTransformer class and overrides the visit_Name method.
  • We set the return value to a Subscript node with the name data and the contents of the slice as a Constant node with the value as the variable name. node.id represents the raw string of the variable name.
  • We can observe before the transformation that a Name node is printed in the tree.
  • We call the constructor of our TransformName class to create an object of that class and use the visit method, which invokes the visit_Name method to transform the Name node to a Subscript node.
  • The ast.fix_missing_locations() method is called on the new tree to relocate all the nodes in the tree, since our TransformName class introduced a new node that was not part of the original tree.
  • We observe that foo is now transformed to data['foo'].

Free Resources

License: Creative Commons-Attribution-ShareAlike 4.0 (CC-BY-SA 4.0)