78

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.

wim
  • 338,267
  • 99
  • 616
  • 750
WGH
  • 3,222
  • 2
  • 29
  • 42

2 Answers2

66

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.

JGC
  • 5,725
  • 1
  • 32
  • 30
Yann
  • 3,841
  • 1
  • 22
  • 14
  • 8
    can you please elaborate on this answer with what key and hashkey are? – Tommy Nov 19 '21 at 01:32
  • 5
    For anyone else wanting an answer to [@Tommy's comment](https://stackoverflow.com/questions/30730983/make-lru-cache-ignore-some-of-the-function-arguments/32655449#comment123792060_32655449), I was able to find the description of `key` and `hashkey` in [this section of the `cachetools` docs](https://cachetools.readthedocs.io/en/stable/#cachetools.cached). 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. – homersimpson Feb 16 '22 at 22:11
  • 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 ` ??? – joseville Feb 24 '22 at 18:11
  • 1
    @Tommy, I edited the answer to include more explanation around the syntax needed. Hope this helps. – JGC Mar 24 '22 at 16:44
  • 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) – stam Feb 10 '23 at 17:58
  • 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 object – stam Feb 10 '23 at 18:13
13

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
WGH
  • 3,222
  • 2
  • 29
  • 42
  • I am reviving this very old answer, but how would you use the `lru_cache_ignoring_first_argument` decorator? – Sanandrea Dec 20 '22 at 17:00