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

Question

I am trying to map the usage / reasons of functions and variables in a python package at the function level. There are several modules in which functions / variables are used in other functions, and I would like to create a dictionary that looks something like this:

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

The functions that I have in mind should be defined in the modules of the package.

How do I start with this? I know that I can iterate through the __dict__ package and check the 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 in the current function. How can this be done? Is there something already developed for this kind of work? I think profiling libraries might have to do something similar to this.

Answer 1

The ast module, as suggested in the comments, ultimately works great. Here is the class I created that is used to retrieve 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 

You can use the class by first importing the python package and passing it to the constructor, and then calling the build method as follows:

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

A dictionary will be printed containing a set of keys representing the name of the function, and values ​​that are lists indicating the functions and / or variables that are used in this function. Not every type of ast counts, but it was good enough in my case.

Replication recursively breaks nodes into simpler types until it reaches ast.Name , after which it can retrieve the name of the variable, function, or method that is used in the target function.