Understanding Python's Abstract Syntax Tree (AST) with Practical Examples

Introduction

In Python programming we sometimes need to process source code for formatting, static analysis, or dynamic generation. Python provides a powerful tool – the Abstract Syntax Tree (AST) – which lets us inspect and manipulate the structure of programs.

What is AST?

An abstract syntax tree is a tree‑like representation of source code where each node corresponds to a syntactic construct such as a function definition or an expression. AST can be used to understand, modify, or even generate new code.

Installation and Environment

The ast module is part of Python’s standard library, so no additional installation is required as long as you have a recent Python version.

Example 1: Build AST and Print

import ast
# Define source code
source_code = """
def hello_world():
    print("Hello, world!")
"""
# Build AST
tree = ast.parse(source_code)
# Print AST structure
print(ast.dump(tree, indent=4))

Output shows the full AST structure of the simple function.

Example 2: Visit Function Definitions

import ast
source_code = """
def hello_world():
    print("Hello, world!")
"""
tree = ast.parse(source_code)
for node in ast.walk(tree):
    if isinstance(node, ast.FunctionDef):
        print(f"Function name: {node.name}")

Output: Function name: hello_world

Example 3: Modify AST

import ast
source_code = """
def hello_world():
    print("Hello, world!")
"""
tree = ast.parse(source_code)
for node in ast.walk(tree):
    if isinstance(node, ast.FunctionDef):
        node.name = "greeting"
new_code = compile(tree, filename="", mode="exec")
exec(new_code)

Calling greeting() would now execute the renamed function.

Example 4: Add Import Statement

import ast
from ast import copy_location, fix_missing_locations
source_code = """
def hello_world():
    print("Hello, world!")
"""
tree = ast.parse(source_code)
import_stmt = ast.Import(names=[ast.alias(name='os', asname=None)])
tree.body.insert(0, import_stmt)
fix_missing_locations(tree)
new_code = compile(tree, filename="", mode="exec")
exec(new_code)

After insertion, the module starts with import os .

Example 5: Replace String Literal

import ast
source_code = """
def hello_world():
    print("Hello, world!")
"""
tree = ast.parse(source_code)
for node in ast.walk(tree):
    if isinstance(node, ast.Str):
        node.s = "你好，世界！"
new_code = compile(tree, filename="", mode="exec")
exec(new_code)

Calling hello_world() now prints the Chinese greeting.

Example 6: Delete a Function

import ast
source_code = """
def hello_world():
    print("Hello, world!")

def goodbye_world():
    print("Goodbye, world!")
"""
tree = ast.parse(source_code)
for node in list(tree.body):
    if isinstance(node, ast.FunctionDef) and node.name == "goodbye_world":
        tree.body.remove(node)
new_code = compile(tree, filename="", mode="exec")
exec(new_code)

Only hello_world() remains.

Example 7: Copy a Function

import ast
from ast import copy_location
source_code = """
def hello_world():
    print("Hello, world!")
"""
tree = ast.parse(source_code)
for node in ast.walk(tree):
    if isinstance(node, ast.FunctionDef) and node.name == "hello_world":
        new_func = copy_location(node, node)
        new_func.name = "say_hello"
        tree.body.append(new_func)
new_code = compile(tree, filename="", mode="exec")
exec(new_code)

Calling say_hello() runs a copy of the original function.

Example 8: Add a Parameter to a Function

import ast
source_code = """
def hello_world():
    print("Hello, world!")
"""
tree = ast.parse(source_code)
for node in ast.walk(tree):
    if isinstance(node, ast.FunctionDef) and node.name == "hello_world":
        node.args.args.append(ast.arg(arg='name', annotation=None))
new_code = compile(tree, filename="", mode="exec")
exec(new_code)

Now hello_world('Alice') can be called with a name argument.

Example 9: Add a Default Parameter

import ast
source_code = """
def hello_world(name):
    print(f"Hello, {name}!")
"""
tree = ast.parse(source_code)
for node in ast.walk(tree):
    if isinstance(node, ast.FunctionDef) and node.name == "hello_world":
        node.args.defaults.append(ast.Constant(value='World'))
new_code = compile(tree, filename="", mode="exec")
exec(new_code)

Calling hello_world() now prints Hello, World! .

Example 10: Use NodeTransformer to Replace Strings

import ast
class StringReplacer(ast.NodeTransformer):
    def visit_Str(self, node):
        return ast.Str(s="新的字符串")
source_code = """
def hello_world():
    print("Hello, world!")
"""
tree = ast.parse(source_code)
tree = StringReplacer().visit(tree)
tree = ast.fix_missing_locations(tree)
new_code = compile(tree, filename="", mode="exec")
exec(new_code)

After transformation, the function prints the new string when called.