Manipulating Python files with Abstract Syntax Trees library

Photo by Jason Leem on Unsplash

Brief Intro to AST

Python’s ast (Abstract Syntax Trees) library is a built-in module that allows you to work with the abstract syntax tree of Python code. It provides tools to parse Python source code into an abstract syntax tree, traverse the tree, and analyze or modify the code's structure programmatically.

The abstract syntax tree represents the hierarchical structure of Python code without containing the code itself. It means code elements such as expressions, statements, functions, classes, and more as nodes in a tree, where each node corresponds to a specific Python construct.

The ast module is commonly used for various purposes, including:

1. Code analysis: It allows you to analyze the structure of Python code and extract information about the code’s components, such as variables, functions, and classes.
2. Code modification: You can programmatically modify Python code by manipulating the abstract syntax tree and then converting it back to valid Python code.
3. Code generation: You can generate Python code dynamically by constructing an abstract syntax tree representing the desired code and then converting it to a string.

The ast module is handy for tools like linters, code analyzers, and code transformers that need to work with Python code on a deeper level. Now let us take a look at one such application,

Checking repeated variables with AST

In the following code, the find_repeated_variable_allocations function takes a script as input, parses it using the ast module, and checks for repeated variable allocations in assignments. If any repeated variables are found, it will print their names. Otherwise, it will indicate that no repeated variable allocation was found.

import ast
from pathlib import Path

def find_repeated_variable_allocations(script):
    tree = ast.parse(script)
    variables = set()
    repeated_allocations = set()

    for node in ast.walk(tree):
        if not isinstance(node, ast.Assign):
            continue

        for target in node.targets:
            if not isinstance(target, ast.Name):
                continue

            var_name = target.id
            if var_name in variables and var_name not in repeated_allocations:
                repeated_allocations.add(var_name)
            variables.add(var_name)

    return repeated_allocations

if __main__ == "__name__":
    
    # Read a py file named script.py
    script_to_test = ""
    file_name = 'script.py'
    with file(Path(__file__).resolve().parent + file_name) as file:
        script_to_test = file.read()
    
    repeated_vars = find_repeated_variable_allocations(script_to_test)
    
    if repeated_vars:
        print("Repeated variable allocation found for:", ", ".join(repeated_vars))
    else:
        print("No repeated variable allocation found.")

In this test, the find_repeated_variable_allocations function takes a script as input, parses it using the ast module, and checks for repeated variable allocations in assignments. If any repeated variables are found, it will print their names. Otherwise, it will indicate that no repeated variable allocation was found.

Conclusion

This is a simple application of ast library and you can basically do all sorts of manipulations of your Python code. It can be used in the migration of bigger projects or editing your scripts at scale. Hope you enjoyed this little hack.