如何确定函数(的源代码)是否包含来自特定模块的方法调用?

时间:2018-12-16 00:51:03

标签: python abstract-syntax-tree static-analysis inspect

比方说,我有一堆函数abcde,我想知道它们是否在调用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仅被导入,但从未调用过其中一种方法。)

我无法修改abcde。因此,我认为有可能为此使用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__)

3 个答案:

答案 0 :(得分:2)

这是一项正在进行的工作,但也许会引发一个更好的主意。我正在使用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

答案 1 :(得分:1)

您只需将模拟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的模拟模块,因此可以拦截函数调用。

如果您不想通过引发异常来中断功能,则仍然可以使用相同的方法,方法是在模拟模块中导入原始的random模块(适当地修改sys.path),然后退回原来的功能。

答案 2 :(得分:1)

您可以将random模块替换为模拟对象,从而提供自定义属性访问并因此拦截函数调用。每当函数之一尝试从random导入(实际上)时,它将实际上访问模拟对象。模拟对象还可以设计为上下文管理器,在测试后退回原始的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