_functools.py 1.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344
  1. import functools
  2. from typing import Callable, TypeVar
  3. from typing_extensions import Concatenate, ParamSpec
  4. _P = ParamSpec("_P")
  5. _T = TypeVar("_T")
  6. _C = TypeVar("_C")
  7. # Sentinel used to indicate that cache lookup failed.
  8. _cache_sentinel = object()
  9. def cache_method(
  10. f: Callable[Concatenate[_C, _P], _T],
  11. ) -> Callable[Concatenate[_C, _P], _T]:
  12. """
  13. Like `@functools.cache` but for methods.
  14. `@functools.cache` (and similarly `@functools.lru_cache`) shouldn't be used
  15. on methods because it caches `self`, keeping it alive
  16. forever. `@cache_method` ignores `self` so won't keep `self` alive (assuming
  17. no cycles with `self` in the parameters).
  18. Footgun warning: This decorator completely ignores self's properties so only
  19. use it when you know that self is frozen or won't change in a meaningful
  20. way (such as the wrapped function being pure).
  21. """
  22. cache_name = "_cache_method_" + f.__name__
  23. @functools.wraps(f)
  24. def wrap(self: _C, *args: _P.args, **kwargs: _P.kwargs) -> _T:
  25. assert not kwargs
  26. if not (cache := getattr(self, cache_name, None)):
  27. cache = {}
  28. setattr(self, cache_name, cache)
  29. cached_value = cache.get(args, _cache_sentinel)
  30. if cached_value is not _cache_sentinel:
  31. return cached_value
  32. value = f(self, *args, **kwargs)
  33. cache[args] = value
  34. return value
  35. return wrap