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.
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().
Apply the transformer and repair location metadata:
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.
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:
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
astwhen 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.
A good production workflow is:
- parse original source
- transform the tree
- unparse to text
- parse the generated text again
- run tests or linters
- 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.parseand transform it withNodeTransformer. - Use
ast.fix_missing_locationsafter editing the tree. - Write modified code back with
ast.unparsewhen 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.

