functools.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409
  1. from __future__ import annotations
  2. __all__ = (
  3. "AsyncCacheInfo",
  4. "AsyncCacheParameters",
  5. "AsyncLRUCacheWrapper",
  6. "cache",
  7. "lru_cache",
  8. "reduce",
  9. )
  10. import functools
  11. import sys
  12. from collections import OrderedDict
  13. from collections.abc import (
  14. AsyncIterable,
  15. Awaitable,
  16. Callable,
  17. Coroutine,
  18. Hashable,
  19. Iterable,
  20. )
  21. from functools import update_wrapper
  22. from inspect import iscoroutinefunction
  23. from typing import (
  24. Any,
  25. Generic,
  26. NamedTuple,
  27. TypedDict,
  28. TypeVar,
  29. cast,
  30. final,
  31. overload,
  32. )
  33. from weakref import WeakKeyDictionary
  34. from ._core._eventloop import current_time
  35. from ._core._synchronization import Lock
  36. from .lowlevel import RunVar, checkpoint
  37. if sys.version_info >= (3, 11):
  38. from typing import ParamSpec
  39. else:
  40. from typing_extensions import ParamSpec
  41. T = TypeVar("T")
  42. S = TypeVar("S")
  43. P = ParamSpec("P")
  44. lru_cache_items: RunVar[
  45. WeakKeyDictionary[
  46. AsyncLRUCacheWrapper[Any, Any],
  47. OrderedDict[
  48. Hashable,
  49. tuple[_InitialMissingType, Lock, float | None]
  50. | tuple[Any, None, float | None],
  51. ],
  52. ]
  53. ] = RunVar("lru_cache_items")
  54. class _InitialMissingType:
  55. pass
  56. initial_missing: _InitialMissingType = _InitialMissingType()
  57. class AsyncCacheInfo(NamedTuple):
  58. hits: int
  59. misses: int
  60. maxsize: int | None
  61. currsize: int
  62. ttl: int | None
  63. class AsyncCacheParameters(TypedDict):
  64. maxsize: int | None
  65. typed: bool
  66. always_checkpoint: bool
  67. ttl: int | None
  68. class _LRUMethodWrapper(Generic[T]):
  69. def __init__(self, wrapper: AsyncLRUCacheWrapper[..., T], instance: object):
  70. self.__wrapper = wrapper
  71. self.__instance = instance
  72. def cache_info(self) -> AsyncCacheInfo:
  73. return self.__wrapper.cache_info()
  74. def cache_parameters(self) -> AsyncCacheParameters:
  75. return self.__wrapper.cache_parameters()
  76. def cache_clear(self) -> None:
  77. self.__wrapper.cache_clear()
  78. async def __call__(self, *args: Any, **kwargs: Any) -> T:
  79. if self.__instance is None:
  80. return await self.__wrapper(*args, **kwargs)
  81. return await self.__wrapper(self.__instance, *args, **kwargs)
  82. @final
  83. class AsyncLRUCacheWrapper(Generic[P, T]):
  84. def __init__(
  85. self,
  86. func: Callable[P, Awaitable[T]],
  87. maxsize: int | None,
  88. typed: bool,
  89. always_checkpoint: bool,
  90. ttl: int | None,
  91. ):
  92. self.__wrapped__ = func
  93. self._hits: int = 0
  94. self._misses: int = 0
  95. self._maxsize = max(maxsize, 0) if maxsize is not None else None
  96. self._currsize: int = 0
  97. self._typed = typed
  98. self._always_checkpoint = always_checkpoint
  99. self._ttl = ttl
  100. update_wrapper(self, func)
  101. def cache_info(self) -> AsyncCacheInfo:
  102. return AsyncCacheInfo(
  103. self._hits, self._misses, self._maxsize, self._currsize, self._ttl
  104. )
  105. def cache_parameters(self) -> AsyncCacheParameters:
  106. return {
  107. "maxsize": self._maxsize,
  108. "typed": self._typed,
  109. "always_checkpoint": self._always_checkpoint,
  110. "ttl": self._ttl,
  111. }
  112. def cache_clear(self) -> None:
  113. if cache := lru_cache_items.get(None):
  114. cache.pop(self, None)
  115. self._hits = self._misses = self._currsize = 0
  116. async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T:
  117. # Easy case first: if maxsize == 0, no caching is done
  118. if self._maxsize == 0:
  119. value = await self.__wrapped__(*args, **kwargs)
  120. self._misses += 1
  121. return value
  122. # The key is constructed as a flat tuple to avoid memory overhead
  123. key: tuple[Any, ...] = args
  124. if kwargs:
  125. # initial_missing is used as a separator
  126. key += (initial_missing,) + sum(kwargs.items(), ())
  127. if self._typed:
  128. key += tuple(type(arg) for arg in args)
  129. if kwargs:
  130. key += (initial_missing,) + tuple(type(val) for val in kwargs.values())
  131. try:
  132. cache = lru_cache_items.get()
  133. except LookupError:
  134. cache = WeakKeyDictionary()
  135. lru_cache_items.set(cache)
  136. try:
  137. cache_entry = cache[self]
  138. except KeyError:
  139. cache_entry = cache[self] = OrderedDict()
  140. cached_value: T | _InitialMissingType
  141. try:
  142. cached_value, lock, expires_at = cache_entry[key]
  143. except KeyError:
  144. # We're the first task to call this function
  145. cached_value, lock, expires_at = (
  146. initial_missing,
  147. Lock(fast_acquire=not self._always_checkpoint),
  148. None,
  149. )
  150. cache_entry[key] = cached_value, lock, expires_at
  151. if lock is None:
  152. if expires_at is not None and current_time() >= expires_at:
  153. self._currsize -= 1
  154. cached_value, lock, expires_at = (
  155. initial_missing,
  156. Lock(fast_acquire=not self._always_checkpoint),
  157. None,
  158. )
  159. cache_entry[key] = cached_value, lock, expires_at
  160. else:
  161. # The value was already cached
  162. self._hits += 1
  163. cache_entry.move_to_end(key)
  164. if self._always_checkpoint:
  165. await checkpoint()
  166. return cast(T, cached_value)
  167. async with lock:
  168. # Check if another task filled the cache while we acquired the lock
  169. if (cached_value := cache_entry[key][0]) is initial_missing:
  170. self._misses += 1
  171. if self._maxsize is not None and self._currsize >= self._maxsize:
  172. cache_entry.popitem(last=False)
  173. else:
  174. self._currsize += 1
  175. value = await self.__wrapped__(*args, **kwargs)
  176. expires_at = (
  177. current_time() + self._ttl if self._ttl is not None else None
  178. )
  179. cache_entry[key] = value, None, expires_at
  180. else:
  181. # Another task filled the cache while we were waiting for the lock
  182. self._hits += 1
  183. cache_entry.move_to_end(key)
  184. value = cast(T, cached_value)
  185. return value
  186. def __get__(
  187. self, instance: object, owner: type | None = None
  188. ) -> _LRUMethodWrapper[T]:
  189. wrapper = _LRUMethodWrapper(self, instance)
  190. update_wrapper(wrapper, self.__wrapped__)
  191. return wrapper
  192. class _LRUCacheWrapper(Generic[T]):
  193. def __init__(
  194. self, maxsize: int | None, typed: bool, always_checkpoint: bool, ttl: int | None
  195. ):
  196. self._maxsize = maxsize
  197. self._typed = typed
  198. self._always_checkpoint = always_checkpoint
  199. self._ttl = ttl
  200. @overload
  201. def __call__( # type: ignore[overload-overlap]
  202. self, func: Callable[P, Coroutine[Any, Any, T]], /
  203. ) -> AsyncLRUCacheWrapper[P, T]: ...
  204. @overload
  205. def __call__(
  206. self, func: Callable[..., T], /
  207. ) -> functools._lru_cache_wrapper[T]: ...
  208. def __call__(
  209. self, f: Callable[P, Coroutine[Any, Any, T]] | Callable[..., T], /
  210. ) -> AsyncLRUCacheWrapper[P, T] | functools._lru_cache_wrapper[T]:
  211. if iscoroutinefunction(f):
  212. return AsyncLRUCacheWrapper(
  213. f, self._maxsize, self._typed, self._always_checkpoint, self._ttl
  214. )
  215. return functools.lru_cache(maxsize=self._maxsize, typed=self._typed)(f) # type: ignore[arg-type]
  216. @overload
  217. def cache( # type: ignore[overload-overlap]
  218. func: Callable[P, Coroutine[Any, Any, T]], /
  219. ) -> AsyncLRUCacheWrapper[P, T]: ...
  220. @overload
  221. def cache(func: Callable[..., T], /) -> functools._lru_cache_wrapper[T]: ...
  222. def cache(
  223. func: Callable[..., T] | Callable[P, Coroutine[Any, Any, T]], /
  224. ) -> AsyncLRUCacheWrapper[P, T] | functools._lru_cache_wrapper[T]:
  225. """
  226. A convenient shortcut for :func:`lru_cache` with ``maxsize=None``.
  227. This is the asynchronous equivalent to :func:`functools.cache`.
  228. """
  229. return lru_cache(maxsize=None)(func)
  230. @overload
  231. def lru_cache(
  232. *,
  233. maxsize: int | None = ...,
  234. typed: bool = ...,
  235. always_checkpoint: bool = ...,
  236. ttl: int | None = ...,
  237. ) -> _LRUCacheWrapper[Any]: ...
  238. @overload
  239. def lru_cache( # type: ignore[overload-overlap]
  240. func: Callable[P, Coroutine[Any, Any, T]], /
  241. ) -> AsyncLRUCacheWrapper[P, T]: ...
  242. @overload
  243. def lru_cache(func: Callable[..., T], /) -> functools._lru_cache_wrapper[T]: ...
  244. def lru_cache(
  245. func: Callable[P, Coroutine[Any, Any, T]] | Callable[..., T] | None = None,
  246. /,
  247. *,
  248. maxsize: int | None = 128,
  249. typed: bool = False,
  250. always_checkpoint: bool = False,
  251. ttl: int | None = None,
  252. ) -> (
  253. AsyncLRUCacheWrapper[P, T] | functools._lru_cache_wrapper[T] | _LRUCacheWrapper[Any]
  254. ):
  255. """
  256. An asynchronous version of :func:`functools.lru_cache`.
  257. If a synchronous function is passed, the standard library
  258. :func:`functools.lru_cache` is applied instead.
  259. :param always_checkpoint: if ``True``, every call to the cached function will be
  260. guaranteed to yield control to the event loop at least once
  261. :param ttl: time in seconds after which to invalidate cache entries
  262. .. note:: Caches and locks are managed on a per-event loop basis.
  263. """
  264. if func is None:
  265. return _LRUCacheWrapper[Any](maxsize, typed, always_checkpoint, ttl)
  266. if not callable(func):
  267. raise TypeError("the first argument must be callable")
  268. return _LRUCacheWrapper[T](maxsize, typed, always_checkpoint, ttl)(func)
  269. @overload
  270. async def reduce(
  271. function: Callable[[T, S], Awaitable[T]],
  272. iterable: Iterable[S] | AsyncIterable[S],
  273. /,
  274. initial: T,
  275. ) -> T: ...
  276. @overload
  277. async def reduce(
  278. function: Callable[[T, T], Awaitable[T]],
  279. iterable: Iterable[T] | AsyncIterable[T],
  280. /,
  281. ) -> T: ...
  282. async def reduce( # type: ignore[misc]
  283. function: Callable[[T, T], Awaitable[T]] | Callable[[T, S], Awaitable[T]],
  284. iterable: Iterable[T] | Iterable[S] | AsyncIterable[T] | AsyncIterable[S],
  285. /,
  286. initial: T | _InitialMissingType = initial_missing,
  287. ) -> T:
  288. """
  289. Asynchronous version of :func:`functools.reduce`.
  290. :param function: a coroutine function that takes two arguments: the accumulated
  291. value and the next element from the iterable
  292. :param iterable: an iterable or async iterable
  293. :param initial: the initial value (if missing, the first element of the iterable is
  294. used as the initial value)
  295. """
  296. element: Any
  297. function_called = False
  298. if isinstance(iterable, AsyncIterable):
  299. async_it = iterable.__aiter__()
  300. if initial is initial_missing:
  301. try:
  302. value = cast(T, await async_it.__anext__())
  303. except StopAsyncIteration:
  304. raise TypeError(
  305. "reduce() of empty sequence with no initial value"
  306. ) from None
  307. else:
  308. value = cast(T, initial)
  309. async for element in async_it:
  310. value = await function(value, element)
  311. function_called = True
  312. elif isinstance(iterable, Iterable):
  313. it = iter(iterable)
  314. if initial is initial_missing:
  315. try:
  316. value = cast(T, next(it))
  317. except StopIteration:
  318. raise TypeError(
  319. "reduce() of empty sequence with no initial value"
  320. ) from None
  321. else:
  322. value = cast(T, initial)
  323. for element in it:
  324. value = await function(value, element)
  325. function_called = True
  326. else:
  327. raise TypeError("reduce() argument 2 must be an iterable or async iterable")
  328. # Make sure there is at least one checkpoint, even if an empty iterable and an
  329. # initial value were given
  330. if not function_called:
  331. await checkpoint()
  332. return value