如何在 python 包中查找 python 函数或变量的所有用途

Question

我正在尝试在函数级别绘制 python 包中函数和变量的用途/原因。有几个模块在其他函数中使用了函数/变量，我想创建一个看起来像这样的字典：

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

我所指的功能需要在包的模块中定义。

我该如何开始呢？我知道我可以遍历包 __dict__ 并通过以下方式测试包中定义的函数：

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

但在那之后我需要找到当前函数中使用的函数。如何才能做到这一点？是否已经为这类工作开发了一些东西？我认为分析库可能必须做类似的事情。

Answer 1

cmets 中建议的ast 模块最终运行良好。这是我创建的一个类，用于提取包中定义的函数或变量，这些函数或变量用于每个函数。

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

可以通过首先导入python包并传递给构造函数来使用该类，然后像这样调用build方法：

import package

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

这将打印出一个字典，其中包含一组表示函数名称的键，以及表示函数中使用的函数和/或变量的列表的值。并非每种 ast 类型都被考虑在内，但在我的情况下这已经足够了。

实现递归地将节点分解为更简单的类型，直到达到ast.Name，之后它可以提取目标函数中正在使用的变量、函数或方法的名称。