如何提取python代码文件中使用的函数？

Question

我想创建一个代码文件中使用的所有函数的列表.例如,如果我们在名为"add_random.py"的文件中有以下代码

`

import numpy as np
from numpy import linalg

def foo():
    print np.random.rand(4) + np.random.randn(4)
    print linalg.norm(np.random.rand(4))

`

我想提取以下列表: [numpy.random.rand, np.random.randn, np.linalg.norm, np.random.rand]

该列表包含代码中使用的函数及其实际名称,格式为"module.submodule.function".是否有用python语言构建的东西可以帮助我做到这一点？

Answer 1

您可以使用以下方法提取所有调用表达式：

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

将此与ast解析树一起使用：

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

演示：

>>> tree = ast.parse('''\
... def foo():
...     print np.random.rand(4) + np.random.randn(4)
...     print linalg.norm(np.random.rand(4))
... ''')
>>> cc = CallCollector()
>>> cc.visit(tree)
>>> cc.calls
['np.random.rand', 'np.random.randn', 'linalg.norm']

上面的walker只处理名称和属性；如果需要更复杂的表达式支持，则必须扩展它。

请注意，收集这样的名称并非易事。任何间接都不会被处理。您可以在函数代码中构建字典来调用并动态交换函数对象，而上述静态分析将无法跟踪它。