Adding a method caching decorator to functools

Ideas
1

Problem Statement

functools.cache and functools.lru_cache are well-known for not being totally safe to apply to methods, or at least not applying in the way a naive user might expect.

Because they cache based on all arguments, including self, classes declared as follows have a “buggy” behavior:

class Foo:
    @functools.cache
    def bar(self) -> int:
        return id(self) + 1

The bug being that Foo instances will never be GCed after Foo.bar() is called, since the cache is holding a reference to them.

This is easily solved with bespoke cache, or you can make your own fancy decorator to try to make this idea generic…

But this is all non-obvious for new users, and the solutions are all higher effort than using functools.cache. I would characterize the need as typically “an annoyance” rather than “a serious gap/fault”. The need for caching/memoizing is common enough that there is a good amount of value to be gained by making a generic solution available in the stdlib.

Proposed Solution

I suggest we add a new decorator in functools, functools.cached_method.

cached_method defaults to a cache of unbounded size, but accepts LRU-cache arguments. The rationale for this is that caching on an instance will more typically want to be unbounded than a global function cache, as the cache has a lifetime limited to the lifetime of the object itself.

In terms of implementation, the cache nests everything under a WeakKeyDictionary whose keys are the first argument to the decorated function, namely self.

I don’t think anything additional is needed for this (?). I’d be enthusiastic about contributing a PR, if interested core devs feel supportive.

I found a few old threads on this topic, but none which directly suggested this solution.

If people want to see it trialed on pypi first, I’m happy to do it, but I find it somewhat doubtful – because the gain is so small vs adding your own caches – that many people would want to pull in a library just for this.

19 Likes
2

This sounds like a reasonable addition to functools.

5 Likes
3

In my experience, the major annoyance when it comes to trying to cache methods is that self needs to be hashable. It doesn’t seem that this proposal addresses that pain point - I would personally be more interested in something standardized that stores the cache as a property on the object, for these cases. That said, this does seem a reasonable solution for the problem presented, I’ve just never encountered that problem in particular while I come up against the hashability requirement with some frequency

1 Like
4

TBH I’d be tempted to do this with the cache stored on the instance itself, rather than as a weak dict. It’s easier to reason about the lifetime of objects that way; the rule is simply “when the object expires, its cache disappears”. Probably wouldn’t play well with __slots__ though.

8 Likes
5

I’d like to propose that the name of the decorator and perhaps the title of the whole discussion should indicate that the main topic is per-instance caching. We don’t want to suggest that @method_cache for methods is what @cache is for functions.

It is more complicated that that. From a little bit broader perspective there are two @cache issues affecting the self argument: the discussed instance lifetime (and GC at its end) and the missing ability to exclude an argument from caching - and that argument is often self, but can be also something else. Many (but not all) problems can be solved by factoring out a core function and applying a cache decorator to that one, so its partially a documentation issue (mini-HOWTO?) as well.

2 Likes
6

Yeah, I have a similar instinct and see the same problem case – if you’re trying to eke out all of the performance you can get, you’re possibly using slotted classes.

This also relates to supporting unhashable types, as mentioned above. If the cache were stored on the instance, then self doesn’t need to be hashable.

But I also don’t know what attribute name we’d use for this. So although there’s an alternative implementation possible here, it’s not as concrete as the weakref dict.

Personally, I’d like to choose the path I can see through to completion, and leave the other way of doing this as “roll your own”.

I do want to suggest, at least to a new developer, that method_cache is to methods as cache is to bare functions.

In fact, in the implementation I propose self is not really excluded from the cache key. It would be possible, with the weakref dict, to have two instances of a class compare equal and therefore share a cache. I don’t see this as a defect; it’s just a natural consequence of the approach.

You’re right that sometimes methods don’t make use of self – you could turn them into static methods or move them off of your classes to make caching work today. But that’s not material to caching on methods which do access the local instance.

I’m open to being swayed to approach this all differently, but I don’t think I understand what you’re suggesting. What should this be called instead of method_cache?

7

That’s very fair. Every choice has consequences, and it’s better to pick something and run with it than to juggle all of them.

As a means of avoiding the “unhashable self” problem, would it be better to use id(self) as the key, and maintain a separate collection of weak refs? But, again, do what you think is right; I’m curious about all the options but don’t take that as criticism.

1 Like
8

Is that not also true with a cache that’s a weak-keyed dictionary keyed on self as well? When the object expires, the entry in the dictionary for that object is removed, i.e., its cache disappears. Or do you mean something different here?

Also, I don’t think unhashable classes should be a showstopper here - by default class instances hash to their id, so for the common case, things would just work. And functools.cache has the same issue with unhashable classes, so I don’t think we should necessarily expect more from this new decorator.

9

Hmm, does the weakref get disposed of promptly as soon as there are no strong refs, or does it wait around until the GC does a pass? Perhaps I’m misunderstanding. Of course, there’s no inherent guarantee that an object will be disposed of promptly even without that, but it is certainly the normal thing.

10

I believe that method caching should usually include the instance, therefore different instances have their own caches, and it makes attaching the cache to instances seem the more appealing option. This is generally safer as usually different instances of the same class shouldn’t have the same result for the same non-self inputs. (and in cases where this isn’t desired, use a function, not a method)

It’s possible to make this “opt-in compatible with slots” if the name of the attribute where the decorators will put the cache is deterministic and documented or user provided. Users wanting to use this with slotted classes need only include that name in slots.

3 Likes
11

this interacts with slots similarly with what I would suggest, as slotted classes will need to explicitly enable weakref support. (by adding __weakref__ to slots)

Either way here seems fine to me as long as the method the decorator will use is documented so that users can compose with slotted classes.

Weakref finalization is relatively unoticable on performance in the absence of reference cycles, but increases the gc pause time in a noticable way if a large number of objects involved in a cycle requiring gc to collect also have weakrefs to finalize. I don’t think this is a reason to decide against this option on it’s own, but it may be something that leads to some people still DIYing caching if impacted.

12

Is there a way to make the existing cache method work for methods as well? I’d prefer changing the behavior to work where expected instead of a new method.

13

Yes, certainly. But it would totally break backwards compatibility.

1 Like
14

I tend to say “no”, until proven otherwise.

It does work! Which is almost precisely the problem. The implications of having a process-global cache whose keys include objects intended to have limited lifetimes are very clear… Once you know to look.

To avoid putting objects into a persistent global cache, the cache has to have specialized behavior. So you need a way to decide to use that different behavior. If there’s a way of making that decision other than having the caller specify (e.g., by function name), then I’m happy to discuss a concrete counterproposal. But it has to be specific, not vaguely “somehow”.

For my part, I can only think of ways which are fragile and incur significant overhead.

FWIW, I’d tend towards it being documented as _{method name}_cache if we go down this route, and make it possible to override.

But I think the whole path here ends up worse for users. For example, once the cache is an attribute of the instances, it has to be handled WRT pickle. I don’t have any particular opinion about what resolution would be there – probably “do nothing” and allow people to decide whether or not their caches should be pickled? – but it highlights that adding attributes increases the surface area.

The point in favor of putting caches on instances is that they don’t share results. I’m thinking about it and can’t imagine scenarios in which you want to define hashing and you want cached methods, but you don’t want instances with the same hash to share a cache. I can try to contrive one, but I’m not able to think one up which occurs naturally. It would mean you want two instances to look the same to an external cache but not to an internal one. Isn’t that very rare?

Overall, I think “caches on instances” feels a little cleaner, but produces a more complex public API. I’d rather accept some (documented!) messy edge cases in exchange for a simpler API.

It’s also possible to support both in the future, FWIW, so this is potentially just about what the default tradeoff should be.

I hadn’t realized that there was this interplay! I’ve never combined weakrefs and slotted classes before.

Regarding the whole discussion of when weakrefs are cleaned up, and perf impact, I believe they are cleaned right after GC deletes an object, but I’m not certain. But the weakref docs seem to intentionally leave some wiggle room here, mentioning that some behaviors are specific to the implementation’s GC.

The perf impact would generally be negligible because the proposal isn’t to use any fancy or nontrivial finalizers.

Right now it looks like sentiment is neutral-to-positive. I’ll try to make some time in the next couple of weeks to work up a draft PR. There’s plenty of time for more feedback, of course.

15

What if functool.[lru_]cache gained a weak: Container[int | str] keyword argument? Arguments at the given indices/names will be cached via weakref. This way it can be used for caching other functions where strong references are undesirable.

Linters can warn about caching methods without specifying weak=[0] (or weak=[] to use the current behavior).

class MyClass:
    @cache(weak=[0])
    def my_method(self, x: int) -> bool:
        ...

Needing to specify index and name separately can be confusing, but lru_cache already caches positional and keyword arguments separately, and the introspection required to match them up automatically is relatively expensive and not always possible.

16

YAGNI. I don’t think that level of flexibility is necessary. (If you actually do need it, it’s not too hard to roll your own.) Done this way, there’d need to then also be a cached_method helper function (if you’re trying to cache a method, you aren’t going to stop and think “hey, I might need to use weak refs on the first argument”, but you might very well notice that the next thing after @cache is @cached_method), and that helper would be pretty much the only usage that the feature gets, so… let’s just make the feature that’s needed.

1 Like
17

Outside of any GC concerns, I think @cached_method should be expected to work wherever @cached_property would work which would require working on classes with unhashable instances. I think the cleanest way to do this would be storing the cache in a defined attribute on an instance over any attempts to try to handle unhashable weakrefs.

The main reason I would want this to work on unhashable classes is that while default classes are hashed by id, the default dataclass is unhashable. Making the class hashable requires either making the whole class frozen or setting unsafe_hash=True, neither of which may be desirable.

With regard to the cache being pickled, cached_property already includes the cached result in the pickle, so I wouldn’t be surprised by cached_method having matching behaviour.

One nice thing you could potentially do is make dataclasses aware of @cached_method and have slotted dataclasses automatically create the necessary slots. Only users creating slotted classes directly would then have to deal with the cache names which I think would go a fair way to making it nicer to use.

5 Likes
18

Chiming in to verify that 99% of cases where I want to cache a method on an unhashable class, it’s on a non-frozen dataclass. unsafe_hash=True isn’t permitted in my projects so it’s not a viable workaround for us.

2 Likes
19

I thought that @cache meant “cache forever” and that @lru_cache meant “cache until the entry ages out”. So the @cache example on a method is doing what it is supposed to. That is potentially useful even beyond the lifetime of the instance:

class T(tuple):
    @cache
    def sorted(self):
        return sorted(self)

Since T implements __eq__ and __hash__, the cache would work for all instances of T that are equal even if one of them stops being used. In this case, the long-lived cache is helpful because it saves a redundant computation. One instance might be gone but another one that is equal to it would still benefit.

If a limited lifetime is desired, switch to @lru_cache, then unused objects age out naturally without need for weak referencing.

20

There are some implementations in the answers to this StackOverflow question: https://stackoverflow.com/questions/33672412

First one:

import functools
import weakref

def memoized_method(*lru_args, **lru_kwargs):
    def decorator(func):
        @functools.wraps(func)
        def wrapped_func(self, *args, **kwargs):
            # We're storing the wrapped method inside the instance. If we had
            # a strong reference to self the instance would never die.
            self_weak = weakref.ref(self)
            @functools.wraps(func)
            @functools.lru_cache(*lru_args, **lru_kwargs)
            def cached_method(*args, **kwargs):
                return func(self_weak(), *args, **kwargs)
            setattr(self, func.__name__, cached_method)
            return cached_method(*args, **kwargs)
        return wrapped_func
    return decorator

Second one:

import functools
import weakref

def weak_lru(maxsize=128, typed=False):
    'LRU Cache decorator that keeps a weak reference to "self"'
    def wrapper(func):

        @functools.lru_cache(maxsize, typed)
        def _func(_self, *args, **kwargs):
            return func(_self(), *args, **kwargs)

        @functools.wraps(func)
        def inner(self, *args, **kwargs):
            return _func(weakref.ref(self), *args, **kwargs)

        return inner

    return wrapper

Third one:

from weakref import WeakKeyDictionary, WeakValueDictionary, ref
from functools import lru_cache, wraps

class IdKey:
    def __init__(self, value):
        self._id = id(value)
    def __hash__(self):
        return self._id
    def __eq__(self, other):
        return self._id == other._id
    def __repr__(self):
        return f"<IdKey(_id={self._id})>"

def lru_cache_method(*lru_args, **lru_kwargs):

    def decorator(method):

        instances = WeakValueDictionary()
        methods = WeakKeyDictionary()

        @wraps(method)
        def cached_method(self, *args, **kwargs):
            key = IdKey(self)
            weakly_bound_cached_method = methods.get(key)
            if weakly_bound_cached_method is None:
                # NOTE This prevents `key` from being GCed until `self` is GCed.
                instances[key] = self
                
                # NOTE This makes sure self can be GCed before `bound_cached_method` is GCed,
                # by avoiding a mutual dependency.
                _self = ref(self)

                @wraps(method)
                @lru_cache(*lru_args, **lru_kwargs)
                def weakly_bound_cached_method(*args, **kwargs):
                    return method(_self(), *args, **kwargs)
                
                # NOTE This entry can be GCed as soon as `self` is GCed.
                methods[key] = weakly_bound_cached_method

            return weakly_bound_cached_method(*args, **kwargs)
        
        return cached_method
    
    return decorator

Two other answers suggest that often all that is needed is to add @staticmethod on top of the @cache. If the result depends on the method arguments but not the instance, this is the easiest way to go.

There is also an implementation in the methodtools project on PyPI,