How to find all uses of a python function or variable in a python package How to find all uses of a python function or variable in a python package python-3.x python-3.x

How to find all uses of a python function or variable in a python package


The ast module as suggested in the comments ended up working nicely. Here is a class that I created which is used to extract the functions or variables defined in the package that are used in each function.

import astimport typesimport inspectclass CausalBuilder(ast.NodeVisitor):    def __init__(self, package):        self.forest = []        self.fnames = []        for name, obj in vars(package).items():            if isinstance(obj, types.ModuleType):                with open(obj.__file__) as f:                    text = f.read()                tree = ast.parse(text)                self.forest.append(tree)            elif isinstance(obj, types.FunctionType):                mod, *_ = inspect.getmodule(obj).__name__.split('.')                if mod == package.__name__:                    self.fnames.append(name)        self.causes = {n: [] for n in self.fnames}    def build(self):        for tree in self.forest:            self.visit(tree)        return self.causes    def visit_FunctionDef(self, node):        self.generic_visit(node)        for b in node.body:            if node.name in self.fnames:                self.causes[node.name] += self.extract_cause(b)    def extract_cause(self, node):        nodes = [node]        cause = []        while nodes:            for i, n in enumerate(nodes):                ntype = type(n)                if ntype == ast.Name:                    if n.id in self.fnames:                        cause.append(n.id)                elif ntype in (ast.Assign, ast.AugAssign, ast.Attribute,                               ast.Subscript, ast.Return):                    nodes.append(n.value)                elif ntype in (ast.If, ast.IfExp):                    nodes.append(n.test)                    nodes.extend(n.body)                    nodes.extend(n.orelse)                elif ntype == ast.Compare:                    nodes.append(n.left)                    nodes.extend(n.comparators)                elif ntype == ast.Call:                    nodes.append(n.func)                elif ntype == ast.BinOp:                    nodes.append(n.left)                    nodes.append(n.right)                elif ntype == ast.UnaryOp:                    nodes.append(n.operand)                elif ntype == ast.BoolOp:                    nodes.extend(n.values)                elif ntype == ast.Num:                    pass                else:                    raise TypeError("Node type `{}` not accounted for."                                    .format(ntype))                nodes.pop(nodes.index(n))        return cause

The class can be used by first importing a python package and passing to the constructor, then calling the build method like so:

import packagecb = CausalBuilder(package)print(cb.build())

Which will print out a dictionary containing a set of keys representing the name of a function, and values which are lists indicating the functions and or variables that are used in the function. Not every ast type is accounted for, but this was good enough in my case.

The implementation recursively breaks down nodes into simpler types until it reaches ast.Name after which it can extract the name of the variable, function, or method that is being used within the target function.