python - 在单元测试期间全局禁用 joblib.memory 缓存

Question

我使用 joblib.Memory 模块来缓存多个模块中的某些函数。缓存在模块和类中分别初始化。

模块1:

memory = Memory(location='/cache/')
@memory.cache
def heavy_function(...)
    .....

模块2:

memory = Memory(location='/cache/')
@memory.cache
def heavy_function2(...)
    .....

模块3:

import Module1
import Module2

def heavy_function3(...)
    Module1.heavy_function1(...)
    Module1.heavy_function1(...)
    .....

现在我有一个单元测试脚本，我想在单元测试期间全局禁用缓存的使用，以确保所有内容都正确计算。 这是否可以在不通过手动为每个模块禁用它的情况下实现 Module1.memory.cachedir=None 还是不删除缓存目录？

我当前的解决方案只是手动修补每个内存调用

单元测试1:

from joblib import Memory
import Module1
Module1.memory = Memory(location=None)
...
unittest.run()

单元测试3:

from joblib import Memory
import Module1 # need to import module 1 just to disable its memory
import Module2 # need to import module 2 just to disable its memory
import Modul3
Module1.memory = Memory(location=None)
Module2.memory = Memory(location=None)
...
unittest.run()

我创建的模块越多，我需要的手动修补的内存就越多。我想可能有更好的解决方案。我在下面提出了一种解决方法。

Answer 1

一种解决方法是在运行测试时设置标志或环境变量。然后在初始化内存之前检查这些标志:

模块1

import os
memflag = os.environ.get('UNITTESTING', False)
memory = Memory(location= None if memflag else '/cache/')
@memory.cache
def heavy_function(...)
    .....

单元测试1

os.environ["UNITTESTING"] = '1'
import Module1
.....
unittest.run()
del os.environ["UNITTESTING"]