Python
AST
Code Modification
Source Code Analysis
Abstract Syntax Tree

Parse a .py file, read the AST, modify it, then write back the modified source code

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

Python's ast module lets you parse source code into a structured syntax tree, transform that tree, and then generate code from the modified result. The important limitation is that ast preserves program structure, not original formatting, so the right write-back strategy depends on whether you care only about semantics or also about comments and exact layout.

Parse the File into an AST

The first step is to read the file and build the tree.

python
1import ast
2from pathlib import Path
3
4path = Path("example.py")
5source = path.read_text(encoding="utf-8")
6tree = ast.parse(source, filename=str(path))
7print(type(tree).__name__)

At this point, you can walk the tree, inspect node types, and decide what transformation to apply. Using AST nodes is safer than text replacement because you are editing Python syntax, not raw characters.

Modify the Tree with NodeTransformer

The standard way to rewrite nodes is ast.NodeTransformer. For example, suppose you want to rename a function call from old_api() to new_api().

python
1import ast
2
3
4class RenameCall(ast.NodeTransformer):
5    def visit_Call(self, node):
6        self.generic_visit(node)
7        if isinstance(node.func, ast.Name) and node.func.id == "old_api":
8            node.func.id = "new_api"
9        return node

Apply the transformer and repair location metadata:

python
transformer = RenameCall()
new_tree = transformer.visit(tree)
ast.fix_missing_locations(new_tree)

fix_missing_locations matters because newly created or modified nodes may otherwise have incomplete line and column metadata.

Write the Tree Back to Source

In Python 3.9 and later, ast.unparse can turn the modified tree back into code.

python
new_source = ast.unparse(new_tree)
path.write_text(new_source + "\n", encoding="utf-8")

This is enough for semantic source-to-source transforms, migrations, and small refactors. The generated code may not match the original formatting, but it should represent the modified program faithfully.

Full Example

Here is a complete working script:

python
1import ast
2from pathlib import Path
3
4
5class RenameCall(ast.NodeTransformer):
6    def visit_Call(self, node):
7        self.generic_visit(node)
8        if isinstance(node.func, ast.Name) and node.func.id == "old_api":
9            node.func.id = "new_api"
10        return node
11
12
13path = Path("example.py")
14source = path.read_text(encoding="utf-8")
15tree = ast.parse(source, filename=str(path))
16new_tree = RenameCall().visit(tree)
17ast.fix_missing_locations(new_tree)
18path.write_text(ast.unparse(new_tree) + "\n", encoding="utf-8")

If example.py originally contained value = old_api(), the rewritten file will call new_api() instead.

Preserve Formatting Only with the Right Tool

This is the point many articles blur. The built-in AST is excellent for syntax-aware transformations, but it does not preserve comments, blank lines, or exact formatting. If that fidelity matters, use a concrete syntax tree tool such as LibCST instead of plain ast.

So the decision is:

  • use ast when semantics matter most
  • use LibCST or a similar library when comments and formatting must survive

That distinction prevents a lot of frustration.

Validate the Output Before Overwriting Important Files

For real refactoring work, do not overwrite the source blindly. Parse the output again and optionally run tests or formatters afterward.

python
candidate = ast.unparse(new_tree)
ast.parse(candidate)

A good production workflow is:

  1. parse original source
  2. transform the tree
  3. unparse to text
  4. parse the generated text again
  5. run tests or linters
  6. then write the file

That catches broken rewrites early.

Common Pitfalls

A common mistake is expecting ast.unparse to preserve comments and exact whitespace. It does not. Another is creating new nodes without repairing location metadata, which can confuse later tooling.

Developers also sometimes use text replacement for syntax changes that should be AST-based. That works until a similar string appears in a comment, string literal, or unrelated identifier.

Finally, do not assume every Python environment has ast.unparse. It is available in newer Python releases, so older environments may need a different code-generation library.

Summary

  • Parse Python source with ast.parse and transform it with NodeTransformer.
  • Use ast.fix_missing_locations after editing the tree.
  • Write modified code back with ast.unparse when semantic correctness matters more than original formatting.
  • Use a concrete syntax tree library if comments and layout must be preserved.
  • Re-parse and validate generated code before overwriting important files.

Course illustration
Course illustration

All Rights Reserved.