Replacing Code With Python: A Practical Guide to AST Transformations

Ever wondered how tools like linters, formatters, or even code analyzers actually work? They don't read your code like a human does. They parse it. And once you understand how Python itself parses code, you can start writing scripts that modify or generate code automatically.

Let's walk through a real example. Imagine you have a large codebase with old-style print statements that you want to convert to logging calls. Doing this manually would take hours. With Python's Abstract Syntax Tree (AST) module, you can write a script that does it in seconds.

What Is an AST?

When Python runs import my_script, it first reads the file and converts it into a tree structure. Each node in that tree represents a piece of code: a function definition, an assignment, a loop, a function call. The root is the entire module, and leaves are things like variable names or literal strings.

The key insight: you can modify that tree before Python compiles it into bytecode. That means you can rewrite code in a safe, structured way.

Your First AST Transformation

Start small. Here's how to parse code and inspect its structure:

import ast

code = """
def greet(name):
    print(f"Hello, {name}")
"""

tree = ast.parse(code)
print(ast.dump(tree, indent=2))

Run that. You'll see something like:

Module(
  body=[
    FunctionDef(
      name='greet',
      args=arguments(...),
      body=[
        Expr(
          value=Call(
            func=Name(id='print'),
            args=[...],
            keywords=[]))])])

This is your code as data. You can walk through it, find nodes of a certain type, and change them.

Real Transformation: Replacing print() with logging.info()

Let's write utility that finds every print() call and replaces it with a logging call. Here's the full script:

import ast
import astor  # pip install astor (to turn AST back to code)

class PrintToLogging(ast.NodeTransformer):
    def visit_Call(self, node):
        # Check if this is a call to 'print'
        if isinstance(node.func, ast.Name) and node.func.id == 'print':
            # Build a call to logging.info with same args
            new_func = ast.Attribute(
                value=ast.Name(id='logging', ctx=ast.Load()),
                attr='info',
                ctx=ast.Load()
            )
            return ast.Call(
                func=new_func,
                args=node.args,
                keywords=node.keywords
            )
        return self.generic_visit(node)

# Read the file
with open('old_code.py', 'r') as f:
    source = f.read()

# Parse and transform
tree = ast.parse(source)
transformer = PrintToLogging()
modified_tree = transformer.visit(tree)

# Fix missing node locations (required by astor)
ast.fix_missing_locations(modified_tree)

# Write back
new_source = astor.to_source(modified_tree)
with open('new_code.py', 'w') as f:
    f.write(new_source)

Run this on a file with print("Hello") and print(user_data). It will rewrite those as logging.info("Hello") and logging.info(user_data).

What Else Can You Do?

AST transformations are not limited to simple replacements. You can:

• Add imports automatically – If you replace print with logging.info, your script can also insert import logging at the top of the file.
• Wrap functions with decorators – Find all functions named get_data and wrap them with @cache.
• Remove dead code – Delete assignments to variables that are never used.
• Rename variables – Change all old_name to new_name safely (only where they're defined, not coincidental string matches).

A Word of Caution

AST works on syntax, not semantics. It cannot understand what your code does. It only sees the structure. So:

• Don't try to replace print in code where print is not the built-in function (e.g., if someone has def print(x): return x → your transformation will break that).
• Always test transformations on a copy of your code first.

Going Further

Python's ast module also has ast.NodeVisitor for reading code without modifying it. Use it to count function calls, detect security issues, or generate documentation.

Tools like Pythonskillset use similar techniques under the hood for linting and auto-fixing. Once you understand AST, you stop thinking of code as text and start seeing it as data you can process.

Try writing your own transformer that renames all variables with single-letter names to more descriptive ones. It's a fun challenge and a great way to learn how your Python environment really works.