Python - AST (Abstract Syntax Tree) Manipulation in Python

Abstract Syntax Tree (AST) manipulation is an advanced Python technique that allows you to analyze, modify, and generate Python code programmatically. Instead of working with raw text (source code as strings), AST works with a structured, tree-like representation of code, which makes transformations safer and more precise.

1. What is an Abstract Syntax Tree?

When Python code is executed, it goes through several stages:

  1. Source Code (your .py file)

  2. Parsing into an Abstract Syntax Tree

  3. Compilation into bytecode

  4. Execution by the Python Virtual Machine

An AST represents the logical structure of the code. Each element in the code becomes a node in the tree.

For example, the code:

x = 5 + 3

is converted into a tree structure where:

  • The assignment is a node

  • The variable x is a node

  • The expression 5 + 3 is another subtree with operator and operands

2. Why Use AST Manipulation?

AST manipulation is useful in many advanced scenarios:

  • Static code analysis (linters like pylint)

  • Code formatting tools (like black)

  • Security auditing

  • Automatic code transformation

  • Building compilers or interpreters

  • Detecting bugs or anti-patterns

  • Refactoring large codebases programmatically

3. Python's ast Module

Python provides a built-in module called ast to work with syntax trees.

Parsing Code into AST

import ast

code = "x = 5 + 3"
tree = ast.parse(code)

print(ast.dump(tree, indent=4))

This converts Python code into a structured tree representation.

4. Understanding AST Nodes

Each part of the code is represented by a specific node type.

Some common node types:

  • Module – root node

  • Assign – assignment

  • BinOp – binary operation

  • Name – variable

  • Constant – literal values

  • Expr – expressions

Example breakdown:

x = 5 + 3

  • Assign

    • target: Name(id='x')

    • value: BinOp(left=5, op=Add, right=3)

5. Traversing the AST

You can walk through the AST using:

Method 1: ast.walk()

for node in ast.walk(tree):
    print(type(node).__name__)

Method 2: NodeVisitor (recommended)

class MyVisitor(ast.NodeVisitor):
    def visit_BinOp(self, node):
        print("Binary operation found")
        self.generic_visit(node)

visitor = MyVisitor()
visitor.visit(tree)

This allows you to inspect specific parts of code.

6. Modifying the AST

To modify code, Python provides NodeTransformer.

Example: Replace all additions with multiplications

class ReplaceAdd(ast.NodeTransformer):
    def visit_BinOp(self, node):
        if isinstance(node.op, ast.Add):
            node.op = ast.Mult()
        return self.generic_visit(node)

tree = ast.parse("x = 5 + 3")
new_tree = ReplaceAdd().visit(tree)

print(ast.unparse(new_tree))

Output:

x = 5 * 3

7. Generating Code from AST

After modifying the AST, you can convert it back to Python code using:

ast.unparse(tree)

In older Python versions, external libraries like astor were used.

8. Practical Use Cases

a. Code Refactoring

Automatically rename variables, update deprecated syntax, or restructure code.

b. Linting Tools

Detect bad coding practices without executing the code.

c. Security Checks

Identify unsafe constructs like eval() usage.

d. Code Instrumentation

Inject logging or performance tracking into existing code.

Example: Automatically adding print statements before function calls.

9. AST vs Regex for Code Manipulation

AST is far more reliable than using regular expressions because:

  • It understands code structure

  • It avoids breaking syntax

  • It works across complex nested constructs

Regex-based code modification is error-prone for non-trivial cases.

10. Limitations

  • AST does not preserve formatting (comments, whitespace)

  • Requires understanding of Python grammar

  • Complex transformations can be difficult to implement

11. Advanced Extensions

  • Writing custom compilers using AST

  • Building domain-specific languages (DSLs)

  • Creating transpilers (Python to another language)

  • Integrating with tools like lib2to3, astroid, or typed_ast

Conclusion

AST manipulation provides deep control over Python code at a structural level. It is widely used in professional tools and large-scale systems where automated code analysis and transformation are required. Mastering AST opens the door to building powerful developer tools, compilers, and intelligent code-processing systems.