Abstract Syntax Trees (ASTs) in Action: Build Your Own Python Code Sniffer

#python #tutorial #programming #learning

Have you ever wondered how linters like flake8, pylint, or even powerful tools like black do their magic? Or how compilers seem to understand your code, rewrite it, optimise it, or yell at you when you forget a semicolon?

One word: AST.

Behind the scenes, ASTs — Abstract Syntax Trees — are the unsung heroes powering everything from syntax highlighting in your editor to code intelligence, transpilers, decompilers, and linters.

In this post, we’ll explore the magic of ASTs and use Python’s built-in ast module to build a custom code sniffer. Think of it as rolling your mini-flake8 — but you decide the rules.

What Is an Abstract Syntax Tree?

An AST (Abstract Syntax Tree) is a tree representation of the syntactic structure of your code.

It breaks your source code into chunks — nodes — where each node represents a construct: a class, function, condition, loop, literal, import, and even fine-grained stuff like operators or variable references.

Below are short Python programs and their tree representation

ASTs help computers understand code like humans do — but in a structured, logical way.

Visualizing ASTs

Here are some tools for visualizing ASTs:

🔍 AST Explorer — live visualization for JavaScript, TypeScript, Python, and more.
🐍 Python’s built-in ast module — lets you programmatically parse and walk through code trees.
🌲 Tree-sitter — great for editor integrations and multi-language parsing.
📦 JavaScript libraries like Babel, Espree — used by tools like ESLint.

When you explore your code this way, it’s surprising how even a five-line function can unfold into dozens of lines in AST form. But that’s the granularity these tools work at — and now you can too.

Meet the Python `ast` Module

Python’s ast module gives you everything you need to parse and analyze Python code from the inside out.

Here’s a quick breakdown:

ast.parse(source_code) turns your code into a tree.
Each part of your code becomes an instance of an ast.AST subclass (FunctionDef, ClassDef, Assign, etc.).
You can walk through the tree using ast.NodeVisitor.

By subclassing NodeVisitor, you can write methods like visit_FunctionDef, visit_ClassDef, etc., to act on specific node types.

Think of it as writing a custom rule engine that says: “Hey, every time you see a class, check if it has a docstring.”

Let’s Build a Simple Code Sniffer

Let’s say we want to write a basic static analyzer that does two things:

Ensures every function and class has a docstring.
Flags any function that’s longer than 40 lines.

Here’s what that looks like in code.

`checker.py`

import ast

class Checker(ast.NodeVisitor):
    def visit_ClassDef(self, node):
        if not ast.get_docstring(node):
            print(f"Class '{node.name}' has no docstring.")
        self.generic_visit(node)

    def visit_FunctionDef(self, node):
        start_line = node.lineno
        end_line = node.end_lineno
        length = end_line - start_line + 1

        if length > 40:
            print(f"Function '{node.name}' has more than 40 lines.")
        if not ast.get_docstring(node):
            print(f"Function '{node.name}' has no docstring.")

        self.generic_visit(node)

`sniffer.py`

from checker import Checker
import argparse
import ast

def sniffer():
    parser = argparse.ArgumentParser(description="Custom Python Code Sniffer")
    parser.add_argument('--file', type=str, required=True, help='Python file to analyze')
    args = parser.parse_args()

    with open(args.file) as f:
        source_code = f.read()

    tree = ast.parse(source_code)
    checker = Checker()
    checker.visit(tree)

if __name__ == "__main__":
    sniffer()

`example.py`

def test():
    pass

Usage

python sniffer.py --file example.py 
function 'test' has no docstring.

The sniffer walks through your Python code, traverses the AST, and calls out any functions or classes missing docstrings or going over the line limit.

Summary

The ability to programmatically inspect and manipulate code at this granular level opens up a world of possibilities for developers. This hands-on example merely scratches the surface of what's possible, illustrating how you can craft bespoke tools to enforce coding standards, refactor code, or even generate new code.

Check out a more modular and extensible version of this sniffer on GitHub, where you can create custom checkers, violations and configure rules.