joblib系列：自定义缓存机制参数化

joblib系列：自定义缓存机制参数化

from functools import wraps
from joblib import Memory


# USE_CACHE_FLAG = False  # 定义是否使用缓存的标志
# 使用缓存机制
USE_CACHE_FLAG = True  # 定义是否使用缓存的标志
CACHE_DIR_PATH = "./joblib_cache"
VERBOSE_FLAG = 0

def cache_if_needed(
    use_cache=True, 
    cache_dir_path='./joblib_cache',
    verbose_flag=0
):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            if use_cache:
                joblib_memory = Memory(
                    location=cache_dir_path, verbose=verbose_flag
                )
                return joblib_memory.cache(func)(*args, **kwargs)
            else:
                return func(*args, **kwargs)
        return wrapper

    return decorator

@cache_if_needed(
    use_cache=USE_CACHE_FLAG, 
    cache_dir_path=CACHE_DIR_PATH,
    verbose_flag=VERBOSE_FLAG
)
def expensive_computation(a, b):
    print("Computing expensive_computation...")
    return a * b + a / b

if __name__ == '__main__':
    
    output_with_cache = expensive_computation(1, 2)
    print(f"==>> output_with_cache: {output_with_cache}")

    # # 不使用缓存机制, 设置USE_CACHE_FLAG=False
    # output_without_cache = expensive_computation(3, 4)
    # print(f"==>> output_without_cache: {output_without_cache}")