Transform Python Code Using AST Nodes

Python AST: Write Code That Writes Code

Ever found yourself doing the same boring code transformations over and over? Renaming variables, switching function calls, or generating boilerplate? There's a smarter way, and it involves Python's not-so-secret weapon: the Abstract Syntax Tree (AST).

What's the big deal?

Most Python developers never touch ASTs directly. We just write code, run it, and move on. But here's the thing - the AST is literally the blueprint of your code. It's the structured representation Python creates before executing anything. And once you learn to read and manipulate that blueprint, you unlock a whole new level of automation.

Think about it this way: When you run import this and see the Zen of Python, somewhere behind the scenes, Python had to parse that code into an AST before doing anything with it. You can intercept that process, modify the blueprint, and change what your code actually does.

Real talk: When would you actually use this?

I've been in situations at PythonSkillset where we needed to analyze thousands of Python files for security vulnerabilities. Manually reading through them? Not feasible. Regex? It breaks on nested structures. But visiting every node in the AST tree? That's how we caught things like unsafe eval() calls in template engines.

Here's a practical example. Let's say you want to find all print() calls in your project and wrap them with timestamp logging:

import ast
import sys

class PrintToLogger(ast.NodeTransformer):
    def visit_Call(self, node):
        if isinstance(node.func, ast.Name) and node.func.id == 'print':
            # Replace with: logger.info(args)
            new_call = ast.Call(
                func=ast.Attribute(value=ast.Name(id='logger'), attr='info'),
                args=node.args,
                keywords=[]
            )
            return ast.fix_missing_locations(new_call)
        return node

def transform_file(filename):
    with open(filename) as f:
        tree = ast.parse(f.read(), filename)

    transformer = PrintToLogger()
    new_tree = transformer.visit(tree)
    ast.fix_missing_locations(new_tree)

    # Compile and execute the modified code
    code = compile(new_tree, filename, 'exec')
    exec(code)

Running this will actually execute your file - but with print() replaced by logger.info(). No regex, no string replacement that breaks on edge cases. Pure structural transformation.

The parts you actually need to know

The AST module comes built into Python (no pip install needed). Here are the tools you'll actually use:

• ast.parse() - Takes your source code string and gives you the tree
• ast.dump() - Shows you the tree structure (great for debugging)
• ast.NodeVisitor - Walk through every node without changing anything
• ast.NodeTransformer - Walk through and modify nodes in place
• compile() - Turn your modified AST back into executable code

The tricky part? Every code construct has a specific node type. An if statement is ast.If, a function call is ast.Call, variables are ast.Name. You'll need to inspect the structure first before transforming.

A real automation script

At PythonSkillset, we maintain a large Django monorepo. When upgrading Django versions, certain function signatures change. Instead of editing 200 files by hand, I wrote a simple AST transformer:

import ast

class UpgradeTransform(ast.NodeTransformer):
    def visit_Call(self, node):
        if (isinstance(node.func, ast.Attribute) and 
            isinstance(node.func.value, ast.Name) and
            node.func.value.id == 'response' and 
            node.func.attr == 'json'):
            # Old: response.json(data)
            # New: response.json(data, encoder=CustomJSONEncoder)
            node.keywords.append(
                ast.keyword(arg='encoder', 
                    value=ast.Attribute(
                        value=ast.Name(id='CustomJSONEncoder'),
                        attr='encode')))
            return ast.fix_missing_locations(node)
        return node

Ran it across the entire codebase. Changes applied everywhere, consistently, in seconds.

But wait, there's a catch

ASTs don't preserve comments or formatting. If you transform code and write it back, all your nicely formatted comments are gone. For large-scale transformations, you're better off using libraries like black or astor that can reproduce the code with formatting.

Also, dynamic code (like eval() and exec()) is dangerous for obvious reasons. When you're modifying ASTs, you're effectively doing the same thing at a lower level. Use it in safe environments - CI pipelines, code generation, static analysis tools.

The mindshift this creates

Once you start thinking in ASTs, you see code differently. That repetitive boilerplate generation? Script it. That migration that needs to change 1000 function calls? One visitor pattern. That linting rule you wish existed? Write your own checker.

Python's AST gives you superpowers without leaving the comfort of regular Python syntax. No external dependencies, no learning a DSL, just pure Python manipulating pure Python.

And next time someone says "write code that writes code," you can nod knowingly. Because you're not generating strings - you're building trees.