Pull all python ast.* classes you use in your codebase! 🐍

gistfile1.txt

#!/usr/bin/env python3

# RUN using: ./find_ast_classes.py ./vyper

import argparse
import ast
import os
import inspect

from ast import NodeVisitor


parser = argparse.ArgumentParser(
        description='',
)
parser.add_argument('paths', help='Paths to check', nargs='+', default=".")
args = parser.parse_args()


class pyASTVisitor(NodeVisitor):

    def __init__(self, file_name, source_code, classes_found):
        self.file_name = file_name
        self.source_code = source_code
        self.classes_found = classes_found

    def visit_Attribute(self, node):
        if isinstance(node.value, ast.Name) and node.value.id == 'ast':
            obj = getattr(ast, node.attr, None)
            if inspect.isclass(obj):
                ast_class_name = f'ast.{node.attr}'
                self.classes_found.setdefault(ast_class_name, set())
                self.classes_found[ast_class_name].add(self.file_name)


def get_ast_class_names(file_name, code, classes_found):
    try:
        nodes = ast.parse(code)
        visitor = pyASTVisitor(file_name, code, classes_found)
        visitor.visit(nodes)
    except Exception:
        print(f'{file_name} failed to parse')


def walk_path(path):
    classes_found = dict()
    for root, dirs, files in os.walk(path, topdown=False):
        for name in files:
            file_name = os.path.join(root, name)
            if '.se.py' not in file_name and file_name.endswith('.py'):
                with open(file_name) as f:
                    get_ast_class_names(file_name, f.read(), classes_found)

    for class_name, paths in classes_found.items():
        print(f'{class_name}:')
        for path in sorted(list(paths)):
            print(f'    {path}')


if __name__ == '__main__':
    for path in args.paths:
        walk_path(path)