在单元测试期间全局禁用 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"]

Answer 2

要全局禁用 joblib.memory 缓存，我会考虑调用 register_store_backend 来覆盖默认的 'local' FileSystemStoreBackend 和 DummyStoreBackend 什么都不做。

类似于以下内容，注意 DummyStoreBackend 是从 Memory 的单元测试中复制的，我还没有测试它是否按预期工作

from joblib.memory import register_store_backend
from joblib._store_backends import StoreBackendBase

class DummyStoreBackend(StoreBackendBase):
    """A dummy store backend that does nothing."""

    def _open_item(self, *args, **kwargs):
        """Open an item on store."""
        "Does nothing"

    def _item_exists(self, location):
        """Check if an item location exists."""
        "Does nothing"

    def _move_item(self, src, dst):
        """Move an item from src to dst in store."""
        "Does nothing"

    def create_location(self, location):
        """Create location on store."""
        "Does nothing"

    def exists(self, obj):
        """Check if an object exists in the store"""
        return False

    def clear_location(self, obj):
        """Clear object on store"""
        "Does nothing"

    def get_items(self):
        """Returns the whole list of items available in cache."""
        return []

    def configure(self, location, *args, **kwargs):
        """Configure the store"""
        "Does nothing"

register_store_backend("local", DummyStoreBackend)

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

Disable joblib.memory caching globally during unittest

python

unit-testing

caching

joblib