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

fin*_*oot 6 python static-analysis abstract-syntax-tree inspect

比方说,我有一堆的功能a,b,c,de我想看看他们称从任何方法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)
Run Code Online (Sandbox Code Playgroud)

我想写一个函数,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
Run Code Online (Sandbox Code Playgroud)

(uses_module(b)False因为random只有进口的,但从来没有它的方法称为一个.)

我不能修改a,b,c,de.所以我认为可以使用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__)
Run Code Online (Sandbox Code Playgroud)

a_g*_*est 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
Run Code Online (Sandbox Code Playgroud)

变量模块名称

指定模块名称的更灵活的方法是通过以下方式实现的:

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
Run Code Online (Sandbox Code Playgroud)