Make @lru_cache ignore some of the function arguments
Asked Answered
Z

2

88

How can I make @functools.lru_cache decorator ignore some of the function arguments with regard to caching key?

For example, I have a function that looks like this:

def find_object(db_handle, query):
    # (omitted code)
    return result

If I apply lru_cache decorator just like that, db_handle will be included in the cache key. As a result, if I try to call the function with the same query, but different db_handle, it will be executed again, which I'd like to avoid. I want lru_cache to consider query argument only.

Zecchino answered 9/6, 2015 at 11:41 Comment(0)
Z
75

With cachetools you can write:

from cachetools import cached
from cachetools.keys import hashkey

from random import randint

@cached(cache={}, key=lambda db_handle, query: hashkey(query))
def find_object(db_handle, query):
    print("processing {0}".format(query))
    return query

queries = list(range(5))
queries.extend(range(5))
for q in queries:
    print("result: {0}".format(find_object(randint(0, 1000), q)))

You will need to install cachetools (pip install cachetools).

The syntax is:

@cached(
    cache={},
    key=lambda <all-function-args>: hashkey(<relevant-args>)
)

Here is another example that includes keyword args:

@cached(
    cache={},
    key=lambda a, b, c=1, d=2: hashkey(a, c)
)
def my_func(a, b, c=1, d=2):
    return a + c

In the example above note that the lambda function input args match the my_func args. You don't have to exactly match the argspec if you don't need to. For example, you can use kwargs to squash out things that aren't needed in the hashkey:

@cached(
    cache={},
    key=lambda a, b, c=1, **kwargs: hashkey(a, c)
)
def my_func(a, b, c=1, d=2, e=3, f=4):
    return a + c

In the above example we don't care about d=, e= and f= args when looking up a cache value, so we can squash them all out with **kwargs.

Zoophilous answered 18/9, 2015 at 15:13 Comment(6)
can you please elaborate on this answer with what key and hashkey are?Caridadcarie
For anyone else wanting an answer to @Tommy's comment, I was able to find the description of key and hashkey in this section of the cachetools docs. Essentially, key is a function which returns a cache key, and hashkey() returns a tuple of its args as an internal cache key and verifies that each arg is hashable.Pastel
What does this print? Does it print: processing 0\n result: 0\n processing 1\n result: 1\n processing 2\n result: 2\n processing 3\n result: 3\n processing 4\n result: 4\n result: 0\n result: 1\n result: 2\n result: 3\n result: 4\n ???Tu
@Tommy, I edited the answer to include more explanation around the syntax needed. Hope this helps.Dire
in general for function def function(a,b): where only b is to be used as a cache key, do @cached(cache={}, key=lambda a,b: hashkey(b)) (syntax is confusing on the first look)Holism
I cannot edit my comment, so 1. yes you must use all the stuff, lambdas etc (I was hoping for smt simple like @cached(key='b') and 2. the cache={} what feels like oddly forcing me to declare default empty parameters is actually "The cache argument specifies a cache object to store previous function arguments and return values" i.e. using dictionary as a cache objectHolism
Z
14

I have at least one very ugly solution. Wrap db_handle in a object that's always equals, and unwrap it inside the function.

It requires a decorator with quite a bit of helper functions, which makes stack trace quite confusing.

class _Equals(object):
    def __init__(self, o):
        self.obj = o

    def __eq__(self, other):
        return True

    def __hash__(self):
        return 0

def lru_cache_ignoring_first_argument(*args, **kwargs):
    lru_decorator = functools.lru_cache(*args, **kwargs)

    def decorator(f):
        @lru_decorator
        def helper(arg1, *args, **kwargs):
            arg1 = arg1.obj
            return f(arg1, *args, **kwargs)

        @functools.wraps(f)
        def function(arg1, *args, **kwargs):
            arg1 = _Equals(arg1)
            return helper(arg1, *args, **kwargs)

        return function

    return decorator
Zecchino answered 9/6, 2015 at 16:53 Comment(2)
I am reviving this very old answer, but how would you use the lru_cache_ignoring_first_argument decorator?Deprecate
@Sanandrea, you just put @lru_cache_ignoring_first_argument()directly above your function.Luncheonette

© 2022 - 2024 — McMap. All rights reserved.