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

How to cache result of function depending on the value returned

我知道 functools.lru_cache and functools.cache(自 Python 3.9 起),但我努力缓存函数的这些参数,这些参数不 return None(或任何其他具体值):

from functools import lru_cache

@lru_cache
def my_fun(link):
    res = fetch_data(link)
    return res
fetch_data 遇到间歇性错误时,

res 为 None。这是我不想缓存结果的时候。

缓存带位置参数的函数

我想我可以使用字典自己实现缓存并仅在 return 值不是 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?.

的答案中找到详细讨论的包装器实现细节(也有参数)

禁止 return 值 - 性能

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