AST in python

History / Edit / PDF / EPUB / BIB /
Created: February 10, 2020 / Updated: November 2, 2024 / Status: finished / 1 min read (~193 words)
Problems Software development Python

I want to analyze a python script to extract something from it. How do I do that?

Python has an abstract syntax tree like most programming language.

You can use the ast module to parse a string that contains the code you want to analyze.

A simple example is as follow. It will read a file defined in the file variable, use ast to parse it, returning a tree that can then be traversed using the visitor pattern. Defining visitors lets you separate the responsibility of each of them, making the code that analyzes code easier to understand.

import ast

class ClassVisitor(ast.NodeVisitor):
    def visit_ClassDef(self, node):
        # Do some logic specific to classes
        self.generic_visit(node)

class FunctionVisitor(ast.NodeVisitor):
    def visit_FunctionDef(self, node):
        # Do some logic specific to functions
        self.generic_visit(node)

visitors = [
    ClassVisitor(),
    FunctionVisitor()
]

with open(file, "r") as f:
    code = f.read()

    tree = ast.parse(code)

    for visitor in visitors:
        visitor.visit(tree)