如何根据返回的值缓存函数的结果

Question

我知道 functools.lru_cache 和 functools.cache（自 Python 3.9 起），但我很难缓存不返回 None（或任何其他特定值）的函数的这些参数：

from functools import lru_cache

@lru_cache
def my_fun(link):
    res = fetch_data(link)
    return res

当fetch_data 遇到间歇性错误时，res 为 None。这是我不希望结果被缓存的时候。

Answer 1

缓存带有位置参数的函数

我想我可以使用字典自己实现缓存并仅在返回值不是None时存储结果：

from functools import wraps

def my_cache(no_cache_result=tuple()):
    if no_cache_result is None:
        no_cache_result = tuple()

    cache = dict()

    def decorator(fun):
        @wraps(fun)
        def wrapper(*args, **kwargs):
            assert len(kwargs) == 0
            if args in cache:
                print('cache: taken from cache')
                return cache[args]
            else:
                res = fun(*args, **kwargs)
                if res not in no_cache_result:
                    print('cache: stored in cache')
                    cache[args] = res
                else:
                    print('cache: NOT stored')
                return res

        return wrapper
    return decorator


@my_cache(no_cache_result=[None])
def my_fun(a):
    print(f'my_fun: called with {a}')
    if a <= 1:
        return a
    else:
        return None


my_fun(0)
my_fun(1)
my_fun(2)

my_fun(0)
my_fun(1)
my_fun(2)

哪些打印（如预期）：

my_fun: called with 0
cache: stored in cache
my_fun: called with 1
cache: stored in cache
my_fun: called with 2
cache: NOT stored
cache: taken from cache
cache: taken from cache
my_fun: called with 2
cache: NOT stored

缓存具有位置和关键字参数的函数

上面的解决方案将可以修饰的函数限制为只有位置参数而不是关键字参数的函数。

以小幅减速为代价，可以通过以下方式进行改进：

def my_cache(no_cache_result=tuple()):
    if no_cache_result is None:
        no_cache_result = tuple()

    cache = dict()

    def decorator(fun):
        @wraps(fun)
        def wrapper(*args, **kwargs):
            _kwargs = tuple(kwargs.items())
            if (args, _kwargs) in cache:
                print('cache: taken from cache')
                return cache[(args, _kwargs)]
            else:
                res = fun(*args, **kwargs)
                if res not in no_cache_result:
                    print('cache: stored in cache')
                    cache[(args, _kwargs)] = res
                else:
                    print('cache: NOT stored')
                return res

        return wrapper
    return decorator

按预期工作：

@my_cache(no_cache_result=[None, ])
def my_fun2(a, b=7):
    print(f'my_fun2: called with {a}, {b}')
    if a <= 1:
        return a
    else:
        return None


my_fun2(0, b=2)
my_fun2(1)
my_fun2(2)

my_fun2(0, b=2)
my_fun2(1)
my_fun2(2)

印刷：

my_fun2: called with 0, 2
cache: stored in cache
my_fun2: called with 1, 7
cache: stored in cache
my_fun2: called with 2, 7
cache: NOT stored
cache: taken from cache
cache: taken from cache
my_fun2: called with 2, 7
cache: NOT stored

它是如何工作的？

装饰器

您可以在Decorators with parameters? 的回答中找到有关包装器实现的详细信息（也包括参数）。

禁止返回值 - 性能

缓存的性能取决于传递的no_cache_result 参数的类型。如果您希望限制对多个返回值的缓存，建议传递一个集合，而不是通常使用的列表，因为 if x in no_cache_result 操作对集合比对列表快得多。