python - 如何查明函数(的源代码)是否包含对特定模块方法的调用？

Question

比方说，我有一堆函数 a、b、c、d 和 e 我想知道他们是否调用了 random 模块中的任何方法:

def a():
    pass

def b():
    import random

def c():
    import random
    random.randint(0, 1)

def d():
    import random as ra
    ra.randint(0, 1)

def e():
    from random import randint as ra
    ra(0, 1)

我想写一个函数 uses_module 这样我就可以期望这些断言通过:

assert uses_module(a) == False
assert uses_module(b) == False
assert uses_module(c) == True
assert uses_module(d) == True
assert uses_module(e) == True

(uses_module(b) 为 False 因为 random 仅被导入，但从未调用其方法之一。)

我无法修改a、b、c、d 和e。所以我认为可以为此使用 ast 并遍历从 inspect.getsource 获得的函数代码。 但我对任何其他建议持开放态度，这只是一个想法。

这就是我使用ast所做的:

def uses_module(function):
    import ast
    import inspect
    nodes = ast.walk(ast.parse(inspect.getsource(function)))
    for node in nodes:
        print(node.__dict__)

Answer 1

这是一项正在进行的工作，但也许它会激发出更好的想法。我正在使用 AST 中的节点类型来尝试断言模块已导入并使用了它提供的某些功能。

我已经添加了一些必要的部分来确定 checker defaultdict 是否属于这种情况，它可以针对某些条件集进行评估，但我没有使用所有键值对来为您的用例建立断言。

def uses_module(function):
    """
    (WIP) assert that a function uses a module
    """
    import ast
    import inspect
    nodes = ast.walk(ast.parse(inspect.getsource(function)))
    checker = defaultdict(set)
    for node in nodes:
        if type(node) in [ast.alias, ast.Import, ast.Name, ast.Attribute]:
            nd = node.__dict__
            if type(node) == ast.alias:
                checker['alias'].add(nd.get('name'))
            if nd.get('name') and nd.get('asname'):
                checker['name'].add(nd.get('name'))
                checker['asname'].add(nd.get('asname'))
            if nd.get('ctx') and nd.get('attr'):
                checker['attr'].add(nd.get('attr'))
            if nd.get('id'):
                checker['id'].add(hex(id(nd.get('ctx'))))
            if nd.get('value') and nd.get('ctx'):
                checker['value'].add(hex(id(nd.get('ctx'))))

    # print(dict(checker)) for debug

    # This check passes your use cases, but probably needs to be expanded
    if checker.get('alias') and checker.get('id'):
        return True
    return False