🟠 Python Internals

Abstract Syntax Trees: Parsing, Walking & Code Generation

📖 Lesson 48 ⏱ 60 min 🧪 5 questions 💻 3 exercises

🎯 Learning Objectives

  • Parse Python source into an AST with ast.parse() and understand the node hierarchy
  • Walk the AST with ast.walk() and ast.NodeVisitor to extract information
  • Transform the AST with ast.NodeTransformer to rewrite code
  • Generate Python source from an AST with ast.unparse()
  • Build a linter / static analyser using AST visitors
  • Understand the key AST node types: Module, FunctionDef, ClassDef, Call, Name, Assign, If, For, expressions vs statements
  • Compile a modified AST back to bytecode and execute it

1 · What is an AST?

After tokenisation, Python's PEG parser builds an Abstract Syntax Tree — a tree of Python objects representing the syntactic structure of code, with all insignificant whitespace and comments stripped away. The AST is the canonical, structured representation that the compiler uses to emit bytecode, and it is the foundation of every tool that needs to reason about Python source programmatically.

The ast module in the standard library gives you full read/write access to the tree. Parsing source into an AST is a single function call:

import ast

source = """
def add(x, y):
    return x + y

result = add(3, 4)
print(result)
"""

tree = ast.parse(source)

# The root is always ast.Module
print(type(tree))                    # <class 'ast.Module'>
print(ast.dump(tree, indent=2))      # full tree dump

# Every node has line/column info (after parse)
for node in ast.walk(tree):
    if hasattr(node, "lineno"):
        print(f"{type(node).__name__:20s} line={node.lineno} col={node.col_offset}")
ast_basics.py

The AST organises nodes into distinct categories:

CategoryKey Node Types
StatementsFunctionDef, AsyncFunctionDef, ClassDef, Return, Assign, AugAssign, AnnAssign, For, While, If, With, Try, Import, ImportFrom, Expr, Delete
ExpressionsBinOp, UnaryOp, BoolOp, Compare, Call, Constant, Name, Attribute, Subscript, List, Tuple, Dict, Set, Lambda, IfExp, ListComp, GeneratorExp, DictComp
Context nodesLoad, Store, Del — appears in Name, Attribute, Subscript to indicate read vs write
Helper nodesarguments, arg, keyword, alias, withitem, ExceptHandler

A mini tree visualiser makes the hierarchy tangible:

import ast

def print_tree(node, indent=0):
    prefix = "  " * indent
    name   = type(node).__name__
    extras = []
    for field, value in ast.iter_fields(node):
        if isinstance(value, (str, int, float, bool, type(None))):
            extras.append(f"{field}={value!r}")
    print(f"{prefix}{name}({', '.join(extras)})")
    for field, value in ast.iter_fields(node):
        if isinstance(value, ast.AST):
            print_tree(value, indent + 1)
        elif isinstance(value, list):
            for item in value:
                if isinstance(item, ast.AST):
                    print_tree(item, indent + 1)

print_tree(ast.parse("x = 1 + 2"))
tree_visualiser.py
Concept: The AST is the interface between the parser and the compiler. Every Python tool that needs to understand code structure — linters (flake8, pylint), type checkers (mypy, pyright), formatters (black), refactoring tools (rope) — works with the AST.

2 · Key Node Types in Depth

Understanding the fields of each node type is essential for writing correct visitors and transformers. Let's examine the most important ones field by field:

import ast

# ── FunctionDef ──
src = '''
def greet(name: str, greeting: str = "Hello") -> str:
    """Return a greeting string."""
    return f"{greeting}, {name}!"
'''
tree = ast.parse(src)
fn = tree.body[0]          # FunctionDef node
print(fn.name)             # "greet"
print(fn.args.args)        # [arg(arg='name', annotation=...), arg(arg='greeting', ...)]
print(fn.args.defaults)    # [Constant(value='Hello')]
print(fn.returns)          # Name(id='str')
print(fn.decorator_list)   # []
print(fn.body[0])          # Expr(value=Constant(value='Return a greeting...'))

# ── Call ──
src2 = "result = func(a, b, key=value)"
call = ast.parse(src2).body[0].value   # Call node
print(type(call))          # <class 'ast.Call'>
print(ast.dump(call.func)) # Name(id='func', ctx=Load())
print([ast.dump(a) for a in call.args])    # [Name(id='a'), Name(id='b')]
print([ast.dump(k) for k in call.keywords])  # [keyword(arg='key', value=Name(id='value'))]

# ── Assign vs AnnAssign vs AugAssign ──
assigns = ast.parse("""
x = 1           # Assign
y: int = 2      # AnnAssign
z += 3          # AugAssign
""").body
print(type(assigns[0]))    # Assign   — targets list + value
print(type(assigns[1]))    # AnnAssign — target + annotation + value
print(type(assigns[2]))    # AugAssign — target + op + value

# ── BinOp ──
binop = ast.parse("a + b * c").body[0].value   # outer +
print(ast.dump(binop, indent=2))
# BinOp(left=Name(id='a'), op=Add(), right=BinOp(left=Name(id='b'), op=Mult(), right=Name(id='c')))
node_types.py

FunctionDef carries name, args (an arguments node with sub-fields for positional, keyword-only, *args, **kwargs), body (list of statements), decorator_list, returns (return annotation), and type_comment. The arguments node further splits into posonlyargs, args, vararg, kwonlyargs, kw_defaults, kwarg, and defaults.

Call nodes have three fields: func (the expression being called), args (positional arguments), and keywords (keyword arguments as keyword nodes). Star-unpacking appears as Starred nodes inside args.

The three assignment forms serve different syntactic patterns: Assign for plain binding (note: targets is a list to support a = b = value), AnnAssign for annotated assignment, and AugAssign for in-place operators.

3 · Walking with ast.NodeVisitor

ast.NodeVisitor is the standard visitor pattern implementation for read-only traversal. You subclass it and define visit_<NodeType> methods — the dispatcher calls the right method based on the node's class name:

import ast
from collections import defaultdict

class FunctionAnalyser(ast.NodeVisitor):
    """Analyses all function definitions in a module."""

    def __init__(self):
        self.functions: list[dict] = []
        self._scope_stack: list[str] = []

    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
        qualname = ".".join(self._scope_stack + [node.name])
        info = {
            "qualname":   qualname,
            "line":       node.lineno,
            "args":       [a.arg for a in node.args.args],
            "has_return_annotation": node.returns is not None,
            "is_async":   False,
            "docstring":  ast.get_docstring(node),
            "decorators": [ast.unparse(d) for d in node.decorator_list],
        }
        self.functions.append(info)
        self._scope_stack.append(node.name)
        self.generic_visit(node)   # recurse into nested functions
        self._scope_stack.pop()

    visit_AsyncFunctionDef = visit_FunctionDef   # same logic for async

    def visit_ClassDef(self, node: ast.ClassDef) -> None:
        self._scope_stack.append(node.name)
        self.generic_visit(node)
        self._scope_stack.pop()


class ImportCollector(ast.NodeVisitor):
    """Collects all imported names from a module."""

    def __init__(self):
        self.imports: list[str] = []

    def visit_Import(self, node: ast.Import) -> None:
        for alias in node.names:
            self.imports.append(alias.asname or alias.name)

    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        for alias in node.names:
            self.imports.append(alias.asname or alias.name)


source = open(__file__).read()   # analyse this very script
tree   = ast.parse(source)

analyser = FunctionAnalyser()
analyser.visit(tree)
for fn in analyser.functions[:3]:
    print(fn)
node_visitor.py

generic_visit(node) calls visit_* on all child nodes. If you override a visit_* method without calling generic_visit, the subtree below that node is not visited. Always call self.generic_visit(node) at the end of visitor methods unless you intentionally want to stop recursion into child nodes.

Tip: You can also use ast.walk(tree) for a flat, breadth-first iteration over all nodes when you don't need visitor dispatch or scope tracking. It's simpler but gives you no control over traversal order or early termination.

4 · Transforming with ast.NodeTransformer

ast.NodeTransformer extends NodeVisitor to allow in-place rewriting. Each visit_* method returns a node — return the original to keep it, return a new node to replace it, or return None to delete it from the tree:

import ast

class ConstantFolder(ast.NodeTransformer):
    """
    Simple constant folding: evaluate BinOp nodes with two Constant children
    at compile time.
    e.g. 2 + 3 → 5, "hello" + " world" → "hello world"
    """
    _OPS = {
        ast.Add:  lambda a, b: a + b,
        ast.Sub:  lambda a, b: a - b,
        ast.Mult: lambda a, b: a * b,
        ast.Div:  lambda a, b: a / b,
    }

    def visit_BinOp(self, node: ast.BinOp) -> ast.AST:
        self.generic_visit(node)   # fold children first (bottom-up)
        op_type = type(node.op)
        if (op_type in self._OPS
                and isinstance(node.left,  ast.Constant)
                and isinstance(node.right, ast.Constant)):
            try:
                result = self._OPS[op_type](node.left.value, node.right.value)
                new_node = ast.Constant(value=result)
                return ast.copy_location(new_node, node)
            except Exception:
                pass
        return node


source = "x = 2 + 3 * 4 + 1"
tree   = ast.parse(source)
print("Before:", ast.unparse(tree))   # x = 2 + 3 * 4 + 1

folded = ConstantFolder().visit(tree)
ast.fix_missing_locations(folded)
print("After: ", ast.unparse(folded))  # x = 15  (all constants folded)

# Compile and execute the transformed tree
code = compile(folded, "<string>", "exec")
ns   = {}
exec(code, ns)
print(ns["x"])   # 15
constant_folder.py

Two helper functions are critical when building new or replacement nodes:

  • ast.copy_location(new_node, old_node) — copies lineno, col_offset, end_lineno, end_col_offset from old to new. Use this whenever you create a replacement node so error messages reference the correct source location.
  • ast.fix_missing_locations(tree) — fills in missing location info by copying from parent nodes. Call this on the entire tree after transformation, before compiling.
Warning: If you forget to call ast.fix_missing_locations before compile(), you'll get a TypeError: required field "lineno" missing from stmt. Always call it as the last step before compilation.

5 · Building a Linter

Combining NodeVisitor with structured violation reporting gives you a working static analyser in under 100 lines. Here's a minimal linter that checks several common style rules:

import ast
import sys
from dataclasses import dataclass, field
from pathlib import Path

@dataclass
class Violation:
    file:    str
    line:    int
    col:     int
    code:    str
    message: str

    def __str__(self):
        return f"{self.file}:{self.line}:{self.col}: {self.code} {self.message}"


class SimpleLinter(ast.NodeVisitor):
    """A minimal linter checking a handful of style rules."""

    def __init__(self, filename: str):
        self.filename   = filename
        self.violations: list[Violation] = []

    def _add(self, node: ast.AST, code: str, msg: str) -> None:
        self.violations.append(Violation(
            self.filename, node.lineno, node.col_offset, code, msg
        ))

    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
        # E001: missing docstring
        if not ast.get_docstring(node):
            self._add(node, "E001", f"function '{node.name}' missing docstring")

        # E002: too many arguments (> 7)
        total_args = len(node.args.args) + len(node.args.posonlyargs) + len(node.args.kwonlyargs)
        if total_args > 7:
            self._add(node, "E002", f"function '{node.name}' has {total_args} args (max 7)")

        # E003: function name not snake_case
        import re
        if not re.match(r'^[a-z_][a-z0-9_]*$', node.name) and not node.name.startswith('_'):
            self._add(node, "E003", f"function '{node.name}' is not snake_case")

        self.generic_visit(node)

    def visit_Import(self, node: ast.Import) -> None:
        # E004: bare star import is not allowed here (just a demo rule)
        pass

    def visit_Compare(self, node: ast.Compare) -> None:
        # E005: `x == None` should be `x is None`
        for op, comp in zip(node.ops, node.comparators):
            if isinstance(op, ast.Eq) and isinstance(comp, ast.Constant) and comp.value is None:
                self._add(node, "E005", "use `is None` instead of `== None`")
        self.generic_visit(node)

    def visit_Assert(self, node: ast.Assert) -> None:
        # E006: assert with a tuple is always True
        if isinstance(node.test, ast.Tuple):
            self._add(node, "E006", "assert with a tuple is always True — did you mean assert a, b?")
        self.generic_visit(node)


def lint_file(path: str) -> list[Violation]:
    source = Path(path).read_text()
    tree   = ast.parse(source, filename=path)
    linter = SimpleLinter(path)
    linter.visit(tree)
    return linter.violations


# Demo
demo_source = '''
def badName(x, y, z, a, b, c, d, e):
    if x == None:
        return 0
    assert (x > 0, "x must be positive")
    return x + y
'''
Path("/tmp/demo.py").write_text(demo_source)
for v in lint_file("/tmp/demo.py"):
    print(v)
simple_linter.py

This linter detects: missing docstrings (E001), excessive arguments (E002), non-snake_case names (E003), equality comparisons with None (E005), and tuple assertions (E006). Real-world linters like flake8 and pylint use exactly the same visitor pattern — they simply have hundreds more rules, plugin architectures, and configuration systems layered on top.

6 · Code Generation with ast.unparse

ast.unparse() (Python 3.9+) converts an AST back to valid Python source code. Combined with programmatic AST construction and transformation, this enables powerful code generation workflows:

import ast, textwrap

def generate_dataclass(name: str, fields: list[tuple[str, str]]) -> str:
    """Generate a @dataclass definition from a list of (field_name, type_hint) pairs."""

    # Build the AST programmatically
    field_nodes = [
        ast.AnnAssign(
            target      = ast.Name(id=fname, ctx=ast.Store()),
            annotation  = ast.Name(id=ftype, ctx=ast.Load()),
            value       = None,
            simple      = 1,
        )
        for fname, ftype in fields
    ]

    class_node = ast.ClassDef(
        name            = name,
        bases           = [],
        keywords        = [],
        body            = field_nodes or [ast.Pass()],
        decorator_list  = [
            ast.Attribute(
                value = ast.Name(id="dataclasses", ctx=ast.Load()),
                attr  = "dataclass",
                ctx   = ast.Load(),
            )
        ],
    )

    module = ast.Module(body=[class_node], type_ignores=[])
    ast.fix_missing_locations(module)

    import_line = "import dataclasses\n\n"
    return import_line + ast.unparse(module)


print(generate_dataclass("Point", [("x", "float"), ("y", "float"), ("label", "str")]))
# import dataclasses
# @dataclasses.dataclass
# class Point:
#     x: float
#     y: float
#     label: str
codegen.py

ast.unparse round-trip fidelity and limitations:

  • Whitespace and comments are LOST — they are not stored in the AST. If you need to preserve formatting, use libcst (a Concrete Syntax Tree library) instead.
  • Operator precedence is faithfully reproduced with explicit parentheses where necessary to preserve semantics.
  • String quoting may change — double quotes may become single quotes or vice versa, since the AST only stores the string value, not the original quoting style.
Tip: For production code generation, consider using ast.unparse as the final step in a pipeline: define templates as source strings, parse them into ASTs, splice in generated nodes, then unparse. This gives you the flexibility of AST manipulation with readable output.

7 · Compiling & Executing Modified ASTs

The full pipeline — parse → transform → compile → exec — lets you treat Python source as data, modify it programmatically, and run the result. The compile() builtin accepts AST nodes directly:

import ast

# ── Pattern: parse → transform → compile → exec ──
source = """
import math

def circle_area(r):
    return math.pi * r ** 2

def circle_perimeter(r):
    return 2 * math.pi * r
"""

tree = ast.parse(source)

# Add a __version__ = "generated" assignment at the top of the module
version_node = ast.Assign(
    targets = [ast.Name(id="__version__", ctx=ast.Store())],
    value   = ast.Constant(value="generated"),
    lineno  = 0, col_offset = 0,
)
tree.body.insert(0, version_node)
ast.fix_missing_locations(tree)

# Compile
code = compile(tree, "<generated>", "exec")

# Execute in a fresh namespace
ns = {}
exec(code, ns)

print(ns["__version__"])                  # "generated"
print(ns["circle_area"](5))               # 78.539...
print(ns["circle_perimeter"](5))          # 31.415...

# ── Pattern: eval mode for expressions ──
expr_tree = ast.parse("x**2 + y**2", mode="eval")
# Inject specific values by rewriting Name nodes
class VarReplacer(ast.NodeTransformer):
    def __init__(self, values: dict):
        self.values = values
    def visit_Name(self, node):
        if node.id in self.values:
            return ast.copy_location(ast.Constant(value=self.values[node.id]), node)
        return node

replaced = VarReplacer({"x": 3, "y": 4}).visit(expr_tree)
ast.fix_missing_locations(replaced)
result = eval(compile(replaced, "<expr>", "eval"))
print(result)   # 25  (3² + 4²)

# ── ast.literal_eval: safe evaluation of literals ──
import ast
safe_data = ast.literal_eval('{"key": [1, 2, 3], "flag": True}')
print(safe_data)   # {'key': [1, 2, 3], 'flag': True}
# Never use eval() for untrusted input — use ast.literal_eval for data structures
compile_exec.py

The compile() function accepts three modes:

ModeAST Root TypeUse Case
"exec"ast.ModuleStatements (module-level code)
"eval"ast.ExpressionSingle expression → returns a value
"single"ast.InteractiveSingle interactive statement (REPL-like)
Warning: ast.literal_eval is the ONLY safe way to evaluate untrusted Python-like data structures. It only accepts literals (str, bytes, int, float, complex, bool, None, list, tuple, dict, set). Any other expression raises ValueError. Never use eval() on untrusted input — it executes arbitrary code.

Real-World AST Applications

The AST is the foundation of Python's entire tooling ecosystem. Here is how major tools use it.

ToolAST techniqueWhat it does
pytestNodeTransformerRewrites assert statements to produce rich failure messages showing sub-expression values
blackParse → normalise → ast.unparseRe-formats code to a canonical style by round-tripping through the AST
mypy / pyrightNodeVisitor + type inferenceWalks the AST resolving types of every expression via a constraint solver
coverage.pyNodeVisitor + sys.settraceFinds all executable lines by walking the AST, then instruments them via trace hooks
flake8 / pylintNodeVisitor pluginsEach check is a visitor that fires on specific node types
numbaParse → custom IR → LLVMCompiles numeric Python functions to native code via its own AST-to-IR pipeline
jinja2 / makoCustom parser + code generationParses template language, generates Python AST, compiles to bytecode
SQLAlchemy ORMExpression objects (AST-like)Builds a query expression tree in Python, then "compiles" it to SQL

pytest's assert rewriting — a mini reproduction

import ast

class AssertRewriter(ast.NodeTransformer):
    """
    Rewrites `assert expr` into code that prints sub-expression values on failure.
    Simplified version of what pytest does.
    """
    def visit_Assert(self, node: ast.Assert) -> ast.AST:
        self.generic_visit(node)

        if not isinstance(node.test, ast.Compare):
            return node   # only rewrite comparisons for this demo

        # Build: if not (left op right): raise AssertionError(f"left={left!r} right={right!r}")
        left  = node.test.left
        right = node.test.comparators[0]

        fmt_msg = ast.JoinedStr(values=[
            ast.Constant(value="assert failed: left="),
            ast.FormattedValue(value=ast.copy_location(ast.parse(ast.unparse(left), mode="eval").body, left),
                               conversion=ord("r"), format_spec=None),
            ast.Constant(value=", right="),
            ast.FormattedValue(value=ast.copy_location(ast.parse(ast.unparse(right), mode="eval").body, right),
                               conversion=ord("r"), format_spec=None),
        ])

        raise_node = ast.Raise(
            exc=ast.Call(
                func=ast.Name(id="AssertionError", ctx=ast.Load()),
                args=[fmt_msg], keywords=[],
            ),
            cause=None,
        )

        if_node = ast.If(
            test=ast.UnaryOp(op=ast.Not(), operand=node.test),
            body=[raise_node],
            orelse=[],
        )
        return ast.copy_location(if_node, node)


src = """
x = 10
y = 20
assert x == y
"""
tree = ast.parse(src)
rewritten = AssertRewriter().visit(tree)
ast.fix_missing_locations(rewritten)
try:
    exec(compile(rewritten, "", "exec"))
except AssertionError as e:
    print(e)   # assert failed: left=10, right=20
assert_rewriter.py

Type Annotation Inspector

A practical visitor that extracts every type annotation from a codebase — the technique used by documentation generators and type stub generators. 

import ast
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

@dataclass
class AnnotationRecord:
    kind:       str         # "function_arg", "return", "variable", "attribute"
    name:       str
    annotation: str
    qualname:   str
    line:       int

class AnnotationExtractor(ast.NodeVisitor):
    def __init__(self, filename: str):
        self.filename  = filename
        self.records:  list[AnnotationRecord] = []
        self._scope:   list[str] = []

    def _qualname(self, name: str) -> str:
        return ".".join(self._scope + [name])

    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
        qualname = self._qualname(node.name)

        # Argument annotations
        for arg in node.args.args + node.args.posonlyargs + node.args.kwonlyargs:
            if arg.annotation:
                self.records.append(AnnotationRecord(
                    kind="function_arg", name=arg.arg,
                    annotation=ast.unparse(arg.annotation),
                    qualname=qualname, line=arg.col_offset,
                ))

        # Return annotation
        if node.returns:
            self.records.append(AnnotationRecord(
                kind="return", name=node.name,
                annotation=ast.unparse(node.returns),
                qualname=qualname, line=node.lineno,
            ))

        self._scope.append(node.name)
        self.generic_visit(node)
        self._scope.pop()

    visit_AsyncFunctionDef = visit_FunctionDef

    def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
        name = ast.unparse(node.target)
        self.records.append(AnnotationRecord(
            kind="variable", name=name,
            annotation=ast.unparse(node.annotation),
            qualname=self._qualname(name),
            line=node.lineno,
        ))
        self.generic_visit(node)

    def visit_ClassDef(self, node: ast.ClassDef) -> None:
        self._scope.append(node.name)
        self.generic_visit(node)
        self._scope.pop()


def extract_annotations(source: str, filename: str = "") -> list[AnnotationRecord]:
    tree      = ast.parse(source, filename)
    extractor = AnnotationExtractor(filename)
    extractor.visit(tree)
    return extractor.records


sample = '''
class Config:
    debug: bool = False
    host: str = "localhost"

    def connect(self, port: int, timeout: float = 5.0) -> bool: ...
'''
for rec in extract_annotations(sample):
    print(f"  {rec.kind:12s}  {rec.qualname:30s}  {rec.annotation}")
annotation_extractor.py

Cyclomatic Complexity Metric

Cyclomatic complexity counts the number of linearly independent paths through a function. It is a proxy for how hard a function is to test and understand. Each if, elif, for, while, except, with, and boolean operator adds 1 to the count. 

import ast
from dataclasses import dataclass

@dataclass
class ComplexityResult:
    name:       str
    qualname:   str
    line:       int
    complexity: int

    def rating(self) -> str:
        if self.complexity <= 5:   return "A (simple)"
        if self.complexity <= 10:  return "B (moderate)"
        if self.complexity <= 20:  return "C (complex)"
        return "D (unmaintainable)"

    def __str__(self):
        return (f"  {self.qualname:40s} CC={self.complexity:3d}  "
                f"{self.rating()}")


class CyclomaticComplexityVisitor(ast.NodeVisitor):
    """Compute McCabe cyclomatic complexity for every function."""

    # These node types each add 1 to the complexity count
    BRANCH_NODES = (
        ast.If, ast.For, ast.AsyncFor, ast.While,
        ast.ExceptHandler, ast.With, ast.AsyncWith,
    )

    def __init__(self):
        self.results: list[ComplexityResult] = []
        self._scope:  list[str] = []

    def _count_complexity(self, node: ast.AST) -> int:
        count = 1   # base complexity
        for child in ast.walk(node):
            if isinstance(child, self.BRANCH_NODES):
                count += 1
            elif isinstance(child, ast.BoolOp):
                count += len(child.values) - 1   # and/or adds branches
        return count

    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
        qualname = ".".join(self._scope + [node.name])
        cc = self._count_complexity(node)
        self.results.append(ComplexityResult(
            name=node.name, qualname=qualname,
            line=node.lineno, complexity=cc,
        ))
        self._scope.append(node.name)
        # Do NOT call generic_visit here — walk() inside _count_complexity
        # already handles nested functions; we visit them manually to get
        # their own entries in results
        for child in ast.iter_child_nodes(node):
            self.visit(child)
        self._scope.pop()

    visit_AsyncFunctionDef = visit_FunctionDef

    def visit_ClassDef(self, node: ast.ClassDef) -> None:
        self._scope.append(node.name)
        self.generic_visit(node)
        self._scope.pop()


def analyse_complexity(source: str, threshold: int = 10) -> None:
    tree    = ast.parse(source)
    visitor = CyclomaticComplexityVisitor()
    visitor.visit(tree)

    print(f"\nCyclomatic Complexity Report  (threshold={threshold})")
    print("=" * 60)
    over = [r for r in visitor.results if r.complexity > threshold]
    for r in sorted(visitor.results, key=lambda x: -x.complexity):
        marker = " ← REFACTOR" if r.complexity > threshold else ""
        print(str(r) + marker)
    if over:
        print(f"\n{len(over)} function(s) exceed threshold {threshold}")


# Test
analyse_complexity(open(__file__).read())
complexity.py

Best Practices

  • Always call ast.fix_missing_locations(tree) after building or modifying AST nodes — every node needs lineno, col_offset, end_lineno, end_col_offset before compile() will accept it.
  • Use ast.copy_location(new_node, old_node) when replacing one node with another — this preserves line number information so tracebacks point to the right line.
  • Call self.generic_visit(node) in NodeTransformer methods before or after your transformation to recurse into children — omitting it stops traversal at that node.
  • Use ast.unparse() for readable debug output and code generation — it is available from Python 3.9+. For older versions use the astunparse third-party package.
  • Prefer ast.literal_eval() over eval() for parsing data literals — it is safe for untrusted input; eval() is never safe for untrusted input.
  • Bottom-up transforms in NodeTransformer: call self.generic_visit(node) BEFORE applying your transformation so child nodes are already transformed when you inspect them.
  • Cache parsed trees when analysing many files — ast.parse() is fast but not free; if you need multiple passes over the same file (lint + type-check + complexity), parse once and share the tree.
  • Test transformers with ast.unparse(before) vs ast.unparse(after) — comparing unparsed strings is the clearest way to assert that a transformation does exactly what you expect.

Exercises

Exercise 1 — Dead Code Detector

Write an AST visitor that detects unreachable statements — code that follows a return, raise, break, or continue at the same indentation level:

  • Visit every function body; for each list of statements, scan for a terminating statement followed by more statements.
  • Report: file, line, and the type of the dead code.
  • Test on a file with intentional dead code (assignment after return, code after raise).
  • Handle nested blocks: dead code inside an if branch is fine if there is an else branch that continues.
💡 Hint 
import ast

TERMINATORS = (ast.Return, ast.Raise, ast.Break, ast.Continue)

def find_dead_code(stmts: list[ast.stmt], filename: str) -> list[str]:
    issues = []
    for i, stmt in enumerate(stmts):
        if isinstance(stmt, TERMINATORS) and i + 1 < len(stmts):
            next_stmt = stmts[i + 1]
            issues.append(
                f"{filename}:{next_stmt.lineno}: unreachable code after "
                f"{type(stmt).__name__}"
            )
    return issues

class DeadCodeVisitor(ast.NodeVisitor):
    def __init__(self, filename):
        self.filename = filename
        self.issues: list[str] = []

    def visit_FunctionDef(self, node):
        self.issues.extend(find_dead_code(node.body, self.filename))
        self.generic_visit(node)

    visit_AsyncFunctionDef = visit_FunctionDef

    def visit_If(self, node):
        self.issues.extend(find_dead_code(node.body, self.filename))
        self.issues.extend(find_dead_code(node.orelse, self.filename))
        self.generic_visit(node)

Exercise 2 — Auto-Timing Decorator Injector

Write a NodeTransformer that automatically wraps every function in a timing decorator — without modifying the original source:

  • Transform every FunctionDef (that isn't already decorated with @timed) to add @timed as the first decorator.
  • The timed decorator should be injected as an import at the top of the module.
  • Use ast.parse + transform + compile + exec to run the transformed module.
  • Verify that calling any function in the transformed module prints its execution time.
💡 Hint 
import ast, time, functools

def timed(fn):
    @functools.wraps(fn)
    def wrapper(*a, **kw):
        t0 = time.perf_counter()
        result = fn(*a, **kw)
        print(f"  {fn.__name__}() took {(time.perf_counter()-t0)*1000:.2f}ms")
        return result
    return wrapper

class TimingInjector(ast.NodeTransformer):
    def visit_FunctionDef(self, node):
        self.generic_visit(node)
        already = any(
            (isinstance(d, ast.Name) and d.id == "timed") or
            (isinstance(d, ast.Attribute) and d.attr == "timed")
            for d in node.decorator_list
        )
        if not already:
            node.decorator_list.insert(0, ast.Name(id="timed", ctx=ast.Load()))
        return node
    visit_AsyncFunctionDef = visit_FunctionDef

src = """
def add(a, b): return a + b
def fib(n): return n if n < 2 else fib(n-1) + fib(n-2)
"""
tree = ast.parse(src)
new_tree = TimingInjector().visit(tree)
ast.fix_missing_locations(new_tree)
ns = {"timed": timed}
exec(compile(new_tree, "", "exec"), ns)
print(ns["add"](3, 4))
print(ns["fib"](10))

Exercise 3 — SQL Query Builder via AST

Build a simple expression compiler that converts Python comparison expressions into SQL WHERE clauses using the AST:

  • Accept a Python expression string like "age > 18 and name == 'Alice'".
  • Parse it with ast.parse(expr, mode="eval").
  • Walk the expression AST, converting: BoolOp(And) → AND, BoolOp(Or) → OR, Compare nodes → SQL comparisons, Name → column name, Constant → literal value.
  • Return the SQL WHERE clause string.
  • Test: "age > 18 and (status == 'active' or role == 'admin')" → age > 18 AND (status = 'active' OR role = 'admin').
💡 Hint 
import ast

OP_MAP = {
    ast.Eq:    "=",  ast.NotEq: "!=",
    ast.Lt:    "<",  ast.LtE:   "<=",
    ast.Gt:    ">",  ast.GtE:   ">=",
    ast.In:    "IN", ast.NotIn: "NOT IN",
    ast.Is:    "IS", ast.IsNot: "IS NOT",
}

def to_sql(node: ast.expr) -> str:
    if isinstance(node, ast.BoolOp):
        op = "AND" if isinstance(node.op, ast.And) else "OR"
        parts = [f"({to_sql(v)})" if isinstance(v, ast.BoolOp) else to_sql(v)
                 for v in node.values]
        return f" {op} ".join(parts)
    elif isinstance(node, ast.Compare):
        left = to_sql(node.left)
        parts = []
        for op, comp in zip(node.ops, node.comparators):
            parts.append(f"{left} {OP_MAP[type(op)]} {to_sql(comp)}")
        return " AND ".join(parts)
    elif isinstance(node, ast.Name):
        return node.id
    elif isinstance(node, ast.Constant):
        return f"'{node.value}'" if isinstance(node.value, str) else str(node.value)
    else:
        raise ValueError(f"Unsupported node: {type(node).__name__}")

def expr_to_sql(expr: str) -> str:
    tree = ast.parse(expr, mode="eval")
    return to_sql(tree.body)

print(expr_to_sql("age > 18 and status == 'active'"))
# age > 18 AND status = 'active'