2 min read

Python AST Module

The ast (Abstract Syntax Tree) module helps Python applications to process trees of the Python abstract syntax grammar. The Abstract Syntax Tree can be modified and compiled back into Python code.

Importing the Module

import ast

Parsing Code

You can parse a string of code into an AST node using ast.parse().

import ast

code = "x = 1 + 2"
tree = ast.parse(code)

print(tree)
# Output: <ast.Module object at ...>

Inspecting the Tree

ast.dump() allows you to inspect the structure of the tree.

import ast

code = "x = 1 + 2"
tree = ast.parse(code)

print(ast.dump(tree, indent=4))

Traversing the Tree

ast.walk()

Recursively yield all descendant nodes in the tree.

import ast

code = """
def hello():
    print("Hello")
"""
tree = ast.parse(code)

for node in ast.walk(tree):
    if isinstance(node, ast.FunctionDef):
        print(f"Function name: {node.name}")

ast.NodeVisitor

A more structured way to traverse the tree is by subclassing ast.NodeVisitor.

import ast

class FuncVisitor(ast.NodeVisitor):
    def visit_FunctionDef(self, node):
        print(f"Found function: {node.name}")
        self.generic_visit(node)

code = """
def add(a, b):
    return a + b
"""
tree = ast.parse(code)
FuncVisitor().visit(tree)

Modifying the Tree

ast.NodeTransformer

Allows you to modify the AST.

import ast

class RewriteAdd(ast.NodeTransformer):
    def visit_BinOp(self, node):
        # Replace addition with subtraction
        if isinstance(node.op, ast.Add):
            return ast.BinOp(left=node.left, op=ast.Sub(), right=node.right)
        return node

code = "print(1 + 2)"
tree = ast.parse(code)
tree = RewriteAdd().visit(tree)

# Compile and execute the modified tree
exec(compile(tree, filename="<ast>", mode="exec"))
# Output: -1 (because 1 - 2 = -1)

programming/python/python