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 astParsing 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)