Posted inProgramming

Abstract Syntax Tree (AST) in Python and TypeScript: Build Custom Linters and Code Transformers

Programming tutorial - IT technology blog
Programming tutorial - IT technology blog

Context & Why

It is 2:15 AM, your pager blares, and the production build crashes. An unescaped legacy SQL call slipped straight past your CI checks. The root cause? A developer formatted the query across four lines with inline comments, completely blinding your regex search. Plain string matching had hit its limit.

Source code is not raw text. It is a nested hierarchy. Compilers and runtimes break your files down into an Abstract Syntax Tree (AST). In an AST, every variable declaration, binary expression, function call, and conditional block becomes a typed node in memory.

Regex linters stumble on nested expressions, multi-line arguments, and code comments. By switching to AST-level analysis, you target exact semantic nodes instead of guessing with string patterns. This allows you to catch syntax-sensitive bugs, enforce strict security rules, and safely run automated migrations (codemods) across hundreds of thousands of lines of code.

Installation

Both Python and JavaScript provide dedicated parsing tooling. Python ships with a native AST module in its standard library, while Babel provides the standard traversal and code generation toolset for TypeScript.

Python Environment Setup

Python includes the built-in ast module. For code generation and tree inspection, install libcst and astor:

# Set up a dedicated virtual environment
python3 -m venv .venv
source .venv/bin/activate

# Install helper libraries for AST manipulation
pip install astor libcst

TypeScript & Node.js Environment Setup

In Node.js, install Babel parser packages along with TypeScript definitions and ts-node:

mkdir ast-transformer && cd ast-transformer
npm init -y

# Install Babel parser, traverser, generator, and AST types
npm install @babel/parser @babel/traverse @babel/generator @babel/types
npm install --save-dev typescript @types/node @types/babel__traverse @types/babel__generator ts-node

Configuration & Implementation

Let’s build two practical tools: a Python security linter that intercepts prohibited function calls, and a TypeScript codemod that replaces deprecated APIs automatically.

1. Building a Custom AST Linter in Python

Assume your engineering team has banned eval() and legacy raw_sql_query() wrappers. Subclassing ast.NodeVisitor allows you to inspect every function call node across a target file:

import ast
import sys

class SecurityLinter(ast.NodeVisitor):
    def __init__(self):
        self.violations = []

    def visit_Call(self, node):
        func_name = None
        if isinstance(node.func, ast.Name):
            func_name = node.func.id
        elif isinstance(node.func, ast.Attribute):
            func_name = node.func.attr

        # Flag prohibited functions
        if func_name in ["eval", "raw_sql_query"]:
            self.violations.append({
                "func": func_name,
                "line": node.lineno,
                "col": node.col_offset
            })
        
        # Continue traversing child nodes
        self.generic_visit(node)

def lint_code(source_code: str):
    tree = ast.parse(source_code)
    linter = SecurityLinter()
    linter.visit(tree)
    return linter.violations

if __name__ == "__main__":
    code_sample = """
# Multi-line query that easily breaks regex checks
result = raw_sql_query(
    "SELECT * FROM users WHERE id = %s",
    user_input
)
eval("2 + 2")
    """
    issues = lint_code(code_sample)
    for issue in issues:
        print(f"[SECURITY VIOLATION] Prohibited call '{issue['func']}' at line {issue['line']}:{issue['col']}")

2. Building an Automated AST Transformer in TypeScript

Manual refactoring across a large repository often introduces regressions. Using Babel’s parser and traverser, we can parse TypeScript files, swap deprecated AST nodes in memory, and regenerate cleanly formatted code.

import * as parser from "@babel/parser";
import _traverse from "@babel/traverse";
import _generate from "@babel/generator";
import * as t from "@babel/types";

// Handle CommonJS/ESM interop
const traverse = (_traverse as any).default || _traverse;
const generate = (_generate as any).default || _generate;

const sourceCode = `
// Target: migrate oldLogger.warn to appLogger.warning
function processPayment(account: string) {
  oldLogger.warn("Payment initiated for: " + account);
  console.log("DEBUG: Processing done");
}
`;

// Step 1: Parse source into AST
const ast = parser.parse(sourceCode, {
  sourceType: "module",
  plugins: ["typescript"],
});

// Step 2: Transform nodes via visitor
traverse(ast, {
  CallExpression(path: any) {
    const callee = path.node.callee;

    // Replace oldLogger.warn(...) with appLogger.warning(...)
    if (
      t.isMemberExpression(callee) &&
      t.isIdentifier(callee.object, { name: "oldLogger" }) &&
      t.isIdentifier(callee.property, { name: "warn" })
    ) {
      path.get("callee").replaceWith(
        t.memberExpression(
          t.identifier("appLogger"),
          t.identifier("warning")
        )
      );
    }

    // Strip console.log statements entirely
    if (
      t.isMemberExpression(callee) &&
      t.isIdentifier(callee.object, { name: "console" }) &&
      t.isIdentifier(callee.property, { name: "log" })
    ) {
      path.remove();
    }
  },
});

// Step 3: Serialize updated AST back to source code
const output = generate(ast, {}, sourceCode);
console.log(output.code);

In a recent framework migration, a codemod structured like this updated 420 TypeScript files—roughly 85,000 lines of code—in 3.8 seconds. It eliminated human error completely, with zero broken syntax tokens or missed multi-line calls.

Verification & Monitoring

Automated AST transformations modify your source code directly. You need strict guardrails to prevent accidental logic removals or broken control flow.

1. Pipeline Integration & Automated Checks

Hook your custom AST linter into pre-commit scripts and CI pipelines. Ensure rule violations fail fast before test suites run:

# Block the build if AST violations are detected
python -m linter_runner src/ || exit 1

2. Idempotency Verification

Codemods must be idempotent. Running your transformation script twice in a row must yield zero file differences on the second run:

# Run 1: Apply transformations
npx ts-node transform.ts src/

# Run 2: Verify git status remains completely clean
npx ts-node transform.ts src/
git diff --exit-code

3. Performance Optimization on Large Repositories

Parsing millions of lines of code into syntax trees increases CPU and memory consumption. Keep these benchmarks in mind:

  • Memory Overhead: A parsed AST typically requires 8x to 12x the memory of the raw source file. For a 50 MB codebase, avoid loading every tree into a single array. Process files in batches using worker pools (such as Node’s worker_threads or Python’s multiprocessing.Pool).
  • Early Node Filtering: Avoid traversing nested function bodies when top-level module statements don’t match your criteria. Calling path.skip() in Babel or returning early in Python cuts traversal overhead by up to 60%.

By shifting static analysis and refactoring from brittle regex patterns to AST traversal, your team can maintain consistent code standards and execute major migrations with complete precision.

Share:
Tags: