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

Question

I'm trying to map out the uses/causes of functions and variables in a python package at the function level. There are several modules where functions/variables are used in other functions, and I'd like to create a dictionary that looks something like:

{'function_name':{'uses': [...functions used in this function...],
                  'causes': [...functions that use this function...]},
 ...
}

The functions that I am referring to need to be defined in modules of the package.

How would I start on this? I know that I can iterate through the package __dict__ and test for functions defined in the package by doing:

import package

import inspect
import types

for name, obj in vars(package).items():
    if isinstance(obj, types.FunctionType):
        module, *_ = inspect.getmodule(obj).__name__.split('.')
        if module == package.__name__:
            # Now that function is obtained need to find usages or functions used within it

But after that I need to find the functions used within the current function. How can this be done? Is there something already developed for this type of work? I think that profiling libraries might have to do something similar to this.

Answer 1

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 ast
import types
import inspect


class 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 package

cb = 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.