在unittest期间全局禁用joblib.memory缓存

Posted 2021-03-28

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

模块1：

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

模块2：

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

单词数：

import Module1
import Module2

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

现在我有一个单元测试脚本，我想在单元测试期间全局禁用缓存，以确保一切都正确计算。这是否可以通过Module1.memory.cachedir=None手动禁用每个模块或不删除cachedir？

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

unittest1：

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

unittest3：

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()

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

答案

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

模块1

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

unittest1

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