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.
import astfrom pprint import pprintclass 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))
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'].
TransformName class that extends the ast.NodeTransformer class and overrides the visit_Name method.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.Name node is printed in the tree.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.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.foo is now transformed to data['foo'].