Introduce a functools.cached_method decorator

resolves #102618

This definition of `cached_method` is based on the discussion on DPO:
   https://discuss.python.org/t/107164

Some tradeoffs need to be made in any version of such a decorator. In
this particular implementation, the choices made are as follows:

- lru_cache will be used under the hood, and `cache_clear()` and
  `cache_info()` will be exposed

- caches will be stored in a separate dict, indexed by `id(self)` --
  meaning instances do not need to be hashable and will not share caches

- weakrefs will be used to delete entries from the cache, so the
  instances must be weak-referencable

- lru_cache itself is not threadsafe, but initialization of the caches
  is threadsafe -- this avoids confusing scenarios in which cache entries
  "disappear"

New documentation is included, marked for 3.16, and a small number of new
tests are added.

Author: sirosen

Doc/library/functools.rst

@@ -122,6 +122,49 @@ The :mod:`!functools` module defines the following functions:
       Python 3.12+ this locking is removed.
 
 
+.. decorator:: cached_method(func)
+
+   Decorator to wrap a method with a bounded or unbounded cache.
+
+   When :func:`cache` or :func:`lru_cache` are used on an instance method, the
+   instance (``self``) will be stored in the cache. As a result, instances cannot be
+   garbage collected until the relevant caches are cleared.
+   This decorator uses :func:`lru_cache`, but it wraps the unbound method to accept
+   a weakref and ensures that caches are cleared when instances are garbage collected.
+
+   This is useful for expensive computations which are consistent with respect to
+   an instance, e.g., those which depend only on immutable attributes.
+
+   Example::
+
+       class DataSet:
+
+           def __init__(self, sequence_of_ints):
+               self._data = tuple(sequence_of_ints)
+
+           @cached_method
+           def shifted(self, shift):
+               return DataSet([value + shift for value in self._data])
+
+   On instances, :func:`cached_method` behaves very similarly to :func:`cache`,
+   providing :func:`cache_info()` and :func:`cache_clear()`.
+
+   The *cached_method* does not prevent all possible race conditions in
+   multi-threaded usage. The function could run more than once on the
+   same instance, with the same inputs, with the latest run setting the cached
+   value. However, initialization of the cached method, which happens lazily on
+   first access, is itself threadsafe.
+
+   This decorator requires that the each instance supports weak references.
+   Some immutable types and slotted classes without ``__weakref__`` as one of
+   the defined slots will encounter errors when the cached method is first used.
+
+   *maxsize* and *typed* are supported as keyword arguments to the decorator,
+   and are passed to the underlying :func:`lru_cache()`.
+
+   .. versionadded:: 3.16
+
+
 .. function:: cmp_to_key(func)
 
    Transform an old-style comparison function to a :term:`key function`.  Used

Lib/functools.py

@@ -16,7 +16,7 @@
 
 from abc import get_cache_token
 from collections import namedtuple
-# import weakref  # Deferred to single_dispatch()
+# import weakref  # Deferred to single_dispatch() and cached_method()
 from operator import itemgetter
 from reprlib import recursive_repr
 from types import FunctionType, GenericAlias, MethodType, MappingProxyType, UnionType
@@ -1183,3 +1183,89 @@ def __get__(self, instance, owner=None):
         return val
 
     __class_getitem__ = classmethod(GenericAlias)
+
+################################################################################
+### cached_method -- a version of lru_cache() which uses `id(self)`
+################################################################################
+
+
+def _cached_method_weakref_callback(cache_dict, id_key):
+    def callback(ref):
+        cache_dict.pop(id_key)
+    return callback
+
+
+def _wrap_unbound_cached_method(ref, unbound_method, maxsize, typed):
+    @lru_cache(maxsize, typed)
+    def wrapped(*args, **kwargs):
+        return unbound_method(ref(), *args, **kwargs)
+    return wrapped
+
+
+class cached_method:
+    """
+    A caching decorator for use on instance methods.
+
+    Using cache or lru_cache on methods is problematic because the instance is put into
+    the cache and cannot be garbage collected until the cache is cleared. This decorator
+    uses a cache based on `id(self)` and a weakref to clear cache entries.
+
+    The instance must be weak-referencable.
+
+    By default, this provides an infinite sized cache similar to functools.cache. Use
+    *maxsize* and *typed* to set these attributes of the underlying LRU cache.
+    """
+    def __init__(self, func=None, /, maxsize=None, typed=False):
+        self.func = None
+        self._maxsize = maxsize
+        self._typed = typed
+        self._function_table = {}
+        # we need a lock when initializing per-instance caches
+        self._cache_init_lock = RLock()
+
+        if func is not None:
+            self.func = func
+            update_wrapper(self, func)
+
+    def __call__(self, func):
+        if self.func is not None:
+            raise TypeError(
+                "Each cached_method decorator can only apply to one function."
+            )
+        self.func = func
+        update_wrapper(self, func)
+        return self
+
+    def __get__(self, instance, owner=None):
+        # similar to singledispatch(), we want to defer use of weakref until/unless it
+        # is needed
+        import weakref
+
+        if instance is None:
+            return self
+
+        instance_id = id(instance)
+
+        # first try to retrieve the cached func without locking (thus avoiding any
+        # unnecessary contention when there is a value), but then retry
+        # under a lock to actually provide safety such that two parallel threads won't
+        # construct distinct caches simultaneously
+        try:
+            ref, cached_func = self._function_table[instance_id]
+        except KeyError:
+            with self._cache_init_lock:
+                try:
+                    ref, cached_func = self._function_table[instance_id]
+                except KeyError:
+                    ref = weakref.ref(
+                        instance,
+                        _cached_method_weakref_callback(
+                            self._function_table, instance_id
+                        ),
+                    )
+                    cached_func = _wrap_unbound_cached_method(
+                        ref, self.func, self._maxsize, self._typed
+                    )
+                    self._function_table[instance_id] = ref, cached_func
+
+        return cached_func

Lib/test/test_functools.py

@@ -3810,5 +3810,55 @@ def prop(self):
         self.assertEqual(t.prop, 1)
 
 
+class CachedValueAdder:
+    def __init__(self, value):
+        self.value = value
+
+    @py_functools.cached_method
+    def add(self, other):
+        return self.value + other
+
+
+class TestCachedMethod(unittest.TestCase):
+    module = py_functools
+
+    def test_cached_usage(self):
+        one = CachedValueAdder(1)
+        self.assertEqual(one.add(2), 3)
+        one.value = 2
+        self.assertEqual(one.add(2), 3)  # still 3, not 4
+        one.add.cache_clear()
+        self.assertEqual(one.add(2), 4)  # now 4
+
+    def test_cache_info(self):
+        one = CachedValueAdder(1)
+        self.assertEqual(one.add.cache_info(),
+            self.module._CacheInfo(hits=0, misses=0, maxsize=None, currsize=0))
+        for _ in range(3):
+            for i in range(10):
+                one.add(i)
+        self.assertEqual(one.add.cache_info(),
+            self.module._CacheInfo(hits=20, misses=10, maxsize=None, currsize=10))
+        one.add.cache_clear()
+        self.assertEqual(one.add.cache_info(),
+            self.module._CacheInfo(hits=0, misses=0, maxsize=None, currsize=0))
+
+    def test_reapplication_causes_type_error(self):
+        with self.assertRaisesRegex(
+            TypeError,
+            r"Each cached_method decorator can only apply to one function\.",
+        ):
+            decorator = py_functools.cached_method()
+
+            class MyObject:
+                @decorator
+                def a(self):
+                    return None
+
+                @decorator
+                def b(self):
+                    return None
+
+
 if __name__ == '__main__':
     unittest.main()