Python AST Module

The AST(Abstract Syntax Tree) module in Python is used for parsing, analyzing, and modifying Python code at the syntactic level. It lets you work with the internal structure of Python code, making it useful for static analysis, code transformation, and metaprogramming.

What is an Abstract Syntax Tree (AST)?

An AST (Abstract Syntax Tree) is a tree representation of the abstract syntactic structure of source code. Each node in the tree represents a construct occurring in the source code, like expressions, statements, function calls, loops, etc.

For example, the Python expression:

x = 1 + 2

Has an AST representation similar to:

Module
 └── Assign
      ├── targets: Name (id='x')
      ├── value: BinOp
            ├── left: Constant (value=1)
            ├── op: Add
            ├── right: Constant (value=2)

How to use the AST Module in Python

The ast module provides functionality to:

• Parse Python code into an AST.
• Examine and alter AST nodes.
• Convert AST back to Python source code.
• Build and run AST.

1. Parsing Python Code into AST

You can convert Python source code into an AST using ast.parse().

Example: Parsing Python Code

import ast

code = "x = 1 + 2"
tree = ast.parse(code)

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

Output:

Module(
    body=[
        Assign(
            targets=[Name(id='x', ctx=Store())],
            value=BinOp(
                left=Constant(value=1),
                op=Add(),
                right=Constant(value=2)
            )
        )
    ],
    type_ignores=[]
)

• Module – Represents the entire Python script.
• Assign → Assignment statement (x = ...).
• BinOp → Binary operation (1 + 2).
• Constant → Integer literals (1 and 2).
• Add → Addition operator.

2. Visiting and Modifying AST Nodes

You can traverse and modify the AST using ast.NodeVisitor and ast.NodeTransformer.

Example: Visiting AST Nodes

You can subclass ast.NodeVisitor to analyze nodes.

class MyVisitor(ast.NodeVisitor):
    def visit_BinOp(self, node):
        print(f"Binary Operation: {ast.dump(node)}")
        self.generic_visit(node)  # Continue visiting child nodes

code = "x = 1 + 2"
tree = ast.parse(code)

visitor = MyVisitor()
visitor.visit(tree)

Output:

Binary Operation: BinOp(left=Constant(value=1), op=Add(), right=Constant(value=2))

Example: Modifying AST Nodes

You can modify nodes using ast.NodeTransformer.

class ReplaceAddWithSub(ast.NodeTransformer):
    def visit_BinOp(self, node):
        if isinstance(node.op, ast.Add):  # Replace "+" with "-"
            node.op = ast.Sub()
        return node

code = "x = 1 + 2"
tree = ast.parse(code)
new_tree = ReplaceAddWithSub().visit(tree)
ast.fix_missing_locations(new_tree)  # Fix source locations

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

Modified AST Output

Module(
    body=[
        Assign(
            targets=[Name(id='x', ctx=Store())],
            value=BinOp(
                left=Constant(value=1),
                op=Sub(),
                right=Constant(value=2)
            )
        )
    ],
    type_ignores=[]
)

Now 1 + 2 is replaced with 1 - 2.

3. Compiling and Executing AST

Once modified, an AST can be compiled into executable Python code using compile().

code = "x = 1 + 2"
tree = ast.parse(code)

compiled_code = compile(tree, filename="<ast>", mode="exec")

exec(compiled_code)
print(x)  # Output: 3

4. Converting AST Back to Source Code

Python does not have a built-in way to convert AST back to source code, but the astunparse module (built into Python 3.9+) can do this.

Example: Unparsing AST

import ast
import astunparse

code = "x = 1 + 2"
tree = ast.parse(code)

print(astunparse.unparse(tree))  # Converts AST back to Python code

Output:

x = 1 + 2

5. Advanced Usage of AST

5.1 Static Code Analysis

You can analyze Python code for certain patterns, like detecting function calls.

class FunctionCallFinder(ast.NodeVisitor):
    def visit_Call(self, node):
        print(f"Function called: {ast.dump(node.func)}")
        self.generic_visit(node)

code = """
def hello():
    print("Hello, World!")
hello()
"""

tree = ast.parse(code)
FunctionCallFinder().visit(tree)

Output:

Function called: Name(id='print', ctx=Load())
Function called: Name(id='hello', ctx=Load())

5.2 Code Transformation (Example: Auto-Logging)

You can insert logging statements before function calls.

class AddLogging(ast.NodeTransformer):
    def visit_FunctionDef(self, node):
        print_stmt = ast.Expr(
            value=ast.Call(
                func=ast.Name(id="print", ctx=ast.Load()),
                args=[ast.Constant(value=f"Calling {node.name}")],
                keywords=[]
            )
        )
        node.body.insert(0, print_stmt)  # Insert logging at function start
        return node

code = """
def greet():
    print("Hello!")

greet()
"""

tree = ast.parse(code)
new_tree = AddLogging().visit(tree)
ast.fix_missing_locations(new_tree)

exec(compile(new_tree, filename="<ast>", mode="exec"))

Output:

Calling greet
Hello!

6. Summary

7. Use Cases

• Static Code Analysis (e.g. linting, security audits).
• Code Transformation (e.g. auto-logging, code optimization).
• Metaprogramming (e.g. generating or modifying code).
• Obfuscation & De-obfuscation (e.g. scrambling code for protection).
• Interpreters & Compilers (e.g. custom Python runtimes).

Conclusion

The ast module provides powerful tools for working with Python code programmatically. Parsing, modifying, and compiling ASTs are methods through which you can analyze and transform Python code in quite advanced ways. It is also widely used in tools such as linters, code formatters, and static analyzers.