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.