get_vars.py

import ast
from typing import override

filename = "ex1.py"
content = open(filename).read()
tree = ast.parse(content)

MODULE_SCOPE = "module"


class MyVisitor(ast.NodeVisitor):
    def __init__(self):
        self.collector = {}
        # we visit a module (file) so general vars will be assigned as module scope
        self._stack = [MODULE_SCOPE]

    def visit_inner(self, node):
        """super.generic_visit without function calls"""
        print(f"visit_inner {ast.dump(node)}")
        for field, value in ast.iter_fields(node):
            if isinstance(node, ast.Call) and field == "func":
                # skip called functions
                continue
            if isinstance(value, list):
                for item in value:
                    if isinstance(item, ast.AST):
                        self.visit(item)
            elif isinstance(value, ast.AST):
                self.visit(value)

    @override
    def generic_visit(self, node):
        # just print the node for debugging
        print(f"entering {ast.dump(node)}")
        # assert we don't pop MODULE_SCOPE by mistake
        assert len(self._stack) > 0
        # when entering a function - append to stack
        if isinstance(node, ast.FunctionDef):
            self._stack.append(node.name)
        elif isinstance(node, ast.Name):
            print("=>", node, node.id)
            current_scope = self._stack[-1]
            if current_scope not in self.collector:
                self.collector[current_scope] = set()
            self.collector[current_scope].add(node.id)
        # continue visit to next nodes
        self.visit_inner(node)
        # when leaving a function - pop from stack
        if isinstance(node, ast.FunctionDef):
            self._stack.pop()


visitor = MyVisitor()
visitor.generic_visit(tree)
print(visitor.collector)