Python AST NodeVisitor and NodeTransformer Guide

Python's standard ast module enables developers to parse source code into an Abstract Syntax Tree (AST) for deep programmatic analysis and manipulation. Within this module, ast.NodeVisitor allows you to traverse and inspect code structures without altering them, making it ideal for linters, metric analyzers, and static analysis tools. In contrast, ast.NodeTransformer subclasses NodeVisitor to allow in-place modification, addition, or removal of nodes, providing a powerful mechanism to automatically refactor, optimize, or rewrite Python code.

Parsing Code into an Abstract Syntax Tree

Before inspecting or transforming code, source text must be converted into tree format using ast.parse(). Each element of the code—such as an assignment, function definition, or binary operation—becomes an instance of an AST node subclass.

import ast

code = """
def calculate(x):
    return x * 2
"""

tree = ast.parse(code)

Inspecting Code with ast.NodeVisitor

ast.NodeVisitor walks through an AST by calling custom visitor methods corresponding to node types.

How NodeVisitor Works

1. Subclassing: Create a class inheriting from ast.NodeVisitor.
2. Define Visitor Methods: Implement methods following the naming pattern visit_<NodeType>(self, node). For example, visit_FunctionDef catches function definitions, while visit_Call catches function calls.
3. Continue Traversal: Call self.generic_visit(node) inside overridden methods to ensure child nodes are also visited. If omitted, traversal stops at that branch.

Example: Inspecting Function Calls

The following example finds and logs all function calls in a snippet:

import ast


class FunctionCallFinder(ast.NodeVisitor):

    def visit_Call(self, node):
        if isinstance(node.func, ast.Name):
            print(
                f"Function called: {node.func.id} at line {node.lineno}"
            )
        self.generic_visit(node)


code_to_check = """
print("Starting")
total = sum([1, 2, 3])
print(total)
"""

tree = ast.parse(code_to_check)
finder = FunctionCallFinder()
finder.visit(tree)

Rewriting Code with ast.NodeTransformer

ast.NodeTransformer extends NodeVisitor with the ability to alter nodes during traversal.

Modification Rules

A visitor method on a NodeTransformer must return one of three options:

• The original node or a replacement node: Swaps the current node with the returned node.
• None: Removes the node from the tree entirely.
• A list of nodes: Replaces the single node with multiple new nodes (useful inside body lists, like function bodies).

Maintaining Tree Integrity

When generating new nodes manually, they often lack source location metadata (lineno, col_offset). Call ast.fix_missing_locations(tree) on the transformed tree to automatically propagate coordinates from parent nodes, preventing runtime compilation errors.

Example: Rewriting Variable Names

The following script modifies code by renaming all instances of a variable named temp to cache:

import ast


class VariableRenamer(ast.NodeTransformer):

    def visit_Name(self, node):
        if node.id == "temp":
            # Return a new node with the updated identifier
            return ast.copy_location(ast.Name(id="cache", ctx=node.ctx), node)
        return node


original_code = """
temp = 42
result = temp * 2
print(temp)
"""

tree = ast.parse(original_code)
transformer = VariableRenamer()
transformed_tree = transformer.visit(tree)
ast.fix_missing_locations(transformed_tree)

Generating Source Code from the Modified Tree

Once ast.NodeTransformer completes its alterations, convert the updated AST back into executable Python source code using ast.unparse() (introduced in Python 3.9):

modified_code = ast.unparse(transformed_tree)
print(modified_code)

Output:

cache = 42
result = cache * 2
print(cache)

Alternatively, compile the tree directly into bytecode using compile(transformed_tree, filename="<ast>", mode="exec") for immediate execution via exec().