时间戳

首页 > 解决方案 > 如何提取 Python 源代码中使用的所有函数和 API 调用？

python - 如何提取 Python 源代码中使用的所有函数和 API 调用？

问题描述

让我们考虑以下 Python 源代码；

def package_data(pkg, roots):
    data = []
    for root in roots:
        for dirname, _, files in os.walk(os.path.join(pkg, root)):
            for fname in files:
                data.append(os.path.relpath(os.path.join(dirname, fname), pkg))

    return {pkg: data}

从这个源代码中，我想提取所有的函数和 API 调用。我发现了一个类似的问题和解决方案。我运行了此处给出的解决方案并生成了输出[os.walk, data.append]。但我正在寻找以下输出[os.walk, os.path.join, data.append, os.path.relpath, os.path.join]。

我在分析以下解决方案代码后了解到，这可以访问第一个括号之前的每个节点并删除其余的东西。

import ast

class CallCollector(ast.NodeVisitor):
    def __init__(self):
        self.calls = []
        self.current = None

    def visit_Call(self, node):
        # new call, trace the function expression
        self.current = ''
        self.visit(node.func)
        self.calls.append(self.current)
        self.current = None

    def generic_visit(self, node):
        if self.current is not None:
            print("warning: {} node in function expression not supported".format(
                  node.__class__.__name__))
        super(CallCollector, self).generic_visit(node)

    # record the func expression 
    def visit_Name(self, node):
        if self.current is None:
            return
        self.current += node.id

    def visit_Attribute(self, node):
        if self.current is None:
            self.generic_visit(node)
        self.visit(node.value)  
        self.current += '.' + node.attr

tree = ast.parse(yoursource)
cc = CallCollector()
cc.visit(tree)
print(cc.calls)

谁能帮我修改这段代码，以便这段代码可以遍历括号内的 API 调用？

注意：这可以在 python 中使用正则表达式来完成。但要找出合适的 API 调用，需要大量的手工劳动。所以，我正在寻找抽象语法树的帮助。

标签： pythonfunctionparsingabstract-syntax-treefunction-call

解决方案

不确定这是否是最好或最简单的解决方案，但至少它确实可以按您的情况工作：

import ast

class CallCollector(ast.NodeVisitor):
    def __init__(self):
        self.calls = []
        self._current = []
        self._in_call = False

    def visit_Call(self, node):
        self._current = []
        self._in_call = True
        self.generic_visit(node)

    def visit_Attribute(self, node):
        if self._in_call:
            self._current.append(node.attr)
        self.generic_visit(node)

    def visit_Name(self, node):
        if self._in_call:
            self._current.append(node.id)
            self.calls.append('.'.join(self._current[::-1]))
            # Reset the state
            self._current = []
            self._in_call = False
        self.generic_visit(node)

给你的例子：

['os.walk', 'os.path.join', 'data.append', 'os.path.relpath', 'os.path.join']

问题是你必须做一个generic_visitin all visits 以确保你正确地走树。之后我还使用了一个列表current来加入（反向）。

我发现不适用于这种方法的一种情况是链式操作，例如：d.setdefault(10, []).append(10).

以防万一您对我如何得出该解决方案感兴趣：

假设一个非常简单的节点访问器实现：

import ast

class CallCollector(ast.NodeVisitor):
    def generic_visit(self, node):
        try:
            print(node, node.id)
        except AttributeError:
            try:
                print(node, node.attr)
            except AttributeError:
                print(node)
        return super().generic_visit(node)

这将打印很多东西，但是如果您查看结果，您会看到一些模式，例如：

...
<_ast.Call object at 0x000001AAEE8FFA58>
<_ast.Attribute object at 0x000001AAEE8FFBE0> walk
<_ast.Name object at 0x000001AAEE8FF518> os
...

和

...
<_ast.Call object at 0x000001AAEE8FF160>
<_ast.Attribute object at 0x000001AAEE8FF588> join
<_ast.Attribute object at 0x000001AAEE8FFC50> path
<_ast.Name object at 0x000001AAEE8FF5C0> os
...

所以首先访问调用节点，然后是属性（如果有的话），最后是名称。因此，当您访问调用节点时，您必须重置状态，将所有属性附加到它并在您点击名称节点时停止。

可以在其中执行此操作，generic_visit但最好在方法中执行此操作visit_Call，...然后generic_visit从这些中调用。

可能需要注意一点：这对于简单的情况很有用，但一旦变得不平凡，这将无法可靠地工作。例如，如果你导入一个子包怎么办？如果将函数绑定到局部变量怎么办？如果你调用结果的getattr结果怎么办？在 Python 中列出由静态分析调用的函数可能是不可能的，因为除了普通问题之外，还有帧破解和动态分配（例如，如果某些导入或调用的函数重新分配了os模块中的名称）。

推荐阅读