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

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

Answer 2

您只需在包含以下代码的本地（测试）目录中放置一个模拟 random.py：

# >= Python 3.7.
def __getattr__(name):
    def mock(*args, **kwargs):
        raise RuntimeError(f'{name}: {args}, {kwargs}')  # For example.
    return mock


# <= Python 3.6.
class Wrapper:
    def __getattr__(self, name):
        def mock(*args, **kwargs):
            raise RuntimeError('{}: {}, {}'.format(name, args, kwargs))  # For example.
        return mock

import sys
sys.modules[__name__] = Wrapper()

然后您只需按如下方式测试您的功能：

def uses_module(func):
    try:
        func()
    except RuntimeError as err:
        print(err)
        return True
    return False

这是可行的，因为它不会导入内置的 random 模块，而是使用模拟 custom attribute access 的模拟模块，因此可以拦截函数调用。

如果您不想通过引发异常来中断函数，您仍然可以使用相同的方法，通过在 mock 模块中导入原始 random 模块（适当修改 sys.path）然后回到原来的功能。

Answer 3

您可以用模拟对象替换 random 模块，提供自定义属性访问，从而拦截函数调用。每当其中一个函数尝试导入（从）random 时，它实际上将访问模拟对象。 mock对象也可以设计成上下文管理器，测试后交还原来的random模块

import sys


class Mock:
    import random
    random = random

    def __enter__(self):
        sys.modules['random'] = self
        self.method_called = False
        return self

    def __exit__(self, *args):
        sys.modules['random'] = self.random

    def __getattr__(self, name):
        def mock(*args, **kwargs):
            self.method_called = True
            return getattr(self.random, name)
        return mock


def uses_module(func):
    with Mock() as m:
        func()
        return m.method_called

可变模块名称

一种更灵活的方式，指定模块的名称，通过以下方式实现：

import importlib
import sys


class Mock:
    def __init__(self, name):
        self.name = name
        self.module = importlib.import_module(name)

    def __enter__(self):
        sys.modules[self.name] = self
        self.method_called = False
        return self

    def __exit__(self, *args):
        sys.modules[self.name] = self.module

    def __getattr__(self, name):
        def mock(*args, **kwargs):
            self.method_called = True
            return getattr(self.module, name)
        return mock


def uses_module(func):
    with Mock('random') as m:
        func()
        return m.method_called

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

How to find out if (the source code of) a function contains a call to a method from a specific module?

python

abstract-syntax-tree

static-analysis

inspect

可变模块名称