Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/anyio/functools.py: 38%

1from __future__ import annotations 

2 

3__all__ = ( 

4 "AsyncCacheInfo", 

5 "AsyncCacheParameters", 

6 "AsyncLRUCacheWrapper", 

7 "cache", 

8 "lru_cache", 

9 "reduce", 

10) 

11 

12import functools 

13from collections import OrderedDict 

14from collections.abc import ( 

15 AsyncIterable, 

16 Awaitable, 

17 Callable, 

18 Coroutine, 

19 Hashable, 

20 Iterable, 

21) 

22from functools import update_wrapper 

23from inspect import iscoroutinefunction 

24from typing import ( 

25 Any, 

26 Generic, 

27 NamedTuple, 

28 ParamSpec, 

29 TypedDict, 

30 TypeVar, 

31 cast, 

32 final, 

33 overload, 

34) 

35from weakref import WeakKeyDictionary 

36 

37from ._core._eventloop import current_time 

38from ._core._synchronization import Lock 

39from .lowlevel import RunVar, checkpoint 

40 

41T = TypeVar("T") 

42S = TypeVar("S") 

43P = ParamSpec("P") 

44lru_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 

55 

56class _InitialMissingType: 

57 pass 

58 

59 

60initial_missing: _InitialMissingType = _InitialMissingType() 

61 

62 

63class AsyncCacheInfo(NamedTuple): 

64 hits: int 

65 misses: int 

66 maxsize: int | None 

67 currsize: int 

68 ttl: int | None 

69 

70 

71class AsyncCacheParameters(TypedDict): 

72 maxsize: int | None 

73 typed: bool 

74 always_checkpoint: bool 

75 ttl: int | None 

76 

77 

78class _LRUMethodWrapper(Generic[T]): 

79 def __init__(self, wrapper: AsyncLRUCacheWrapper[..., T], instance: object): 

80 self.__wrapper = wrapper 

81 self.__instance = instance 

82 

83 def cache_info(self) -> AsyncCacheInfo: 

84 return self.__wrapper.cache_info() 

85 

86 def cache_parameters(self) -> AsyncCacheParameters: 

87 return self.__wrapper.cache_parameters() 

88 

89 def cache_clear(self) -> None: 

90 self.__wrapper.cache_clear() 

91 

92 async def __call__(self, *args: Any, **kwargs: Any) -> T: 

93 if self.__instance is None: 

94 return await self.__wrapper(*args, **kwargs) 

95 

96 return await self.__wrapper(self.__instance, *args, **kwargs) 

97 

98 

99@final 

100class AsyncLRUCacheWrapper(Generic[P, T]): 

101 def __init__( 

102 self, 

103 func: Callable[P, Awaitable[T]], 

104 maxsize: int | None, 

105 typed: bool, 

106 always_checkpoint: bool, 

107 ttl: int | None, 

108 ): 

109 self.__wrapped__ = func 

110 self._hits: int = 0 

111 self._misses: int = 0 

112 self._maxsize = max(maxsize, 0) if maxsize is not None else None 

113 self._currsize: int = 0 

114 self._typed = typed 

115 self._always_checkpoint = always_checkpoint 

116 self._ttl = ttl 

117 update_wrapper(self, func) 

118 

119 def cache_info(self) -> AsyncCacheInfo: 

120 return AsyncCacheInfo( 

121 self._hits, self._misses, self._maxsize, self._currsize, self._ttl 

122 ) 

123 

124 def cache_parameters(self) -> AsyncCacheParameters: 

125 return { 

126 "maxsize": self._maxsize, 

127 "typed": self._typed, 

128 "always_checkpoint": self._always_checkpoint, 

129 "ttl": self._ttl, 

130 } 

131 

132 def cache_clear(self) -> None: 

133 if cache := lru_cache_items.get(None): 

134 cache.pop(self, None) 

135 self._hits = self._misses = self._currsize = 0 

136 

137 async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T: 

138 # Easy case first: if maxsize == 0, no caching is done 

139 if self._maxsize == 0: 

140 value = await self.__wrapped__(*args, **kwargs) 

141 self._misses += 1 

142 return value 

143 

144 # The key is constructed as a flat tuple to avoid memory overhead 

145 key: tuple[Any, ...] = args 

146 if kwargs: 

147 # initial_missing is used as a separator 

148 key += (initial_missing,) + sum(kwargs.items(), ()) 

149 

150 if self._typed: 

151 key += tuple(type(arg) for arg in args) 

152 if kwargs: 

153 key += (initial_missing,) + tuple(type(val) for val in kwargs.values()) 

154 

155 try: 

156 cache = lru_cache_items.get() 

157 except LookupError: 

158 cache = WeakKeyDictionary() 

159 lru_cache_items.set(cache) 

160 

161 try: 

162 cache_entry = cache[self] 

163 except KeyError: 

164 cache_entry = cache[self] = OrderedDict() 

165 

166 cached_value: T | _InitialMissingType 

167 try: 

168 cached_value, lock, expires_at = cache_entry[key] 

169 except KeyError: 

170 # We're the first task to call this function 

171 cached_value, lock, expires_at = ( 

172 initial_missing, 

173 Lock(fast_acquire=not self._always_checkpoint), 

174 None, 

175 ) 

176 cache_entry[key] = cached_value, lock, expires_at 

177 

178 if lock is None: 

179 if expires_at is not None and current_time() >= expires_at: 

180 self._currsize -= 1 

181 cached_value, lock, expires_at = ( 

182 initial_missing, 

183 Lock(fast_acquire=not self._always_checkpoint), 

184 None, 

185 ) 

186 cache_entry[key] = cached_value, lock, expires_at 

187 else: 

188 # The value was already cached 

189 self._hits += 1 

190 cache_entry.move_to_end(key) 

191 if self._always_checkpoint: 

192 await checkpoint() 

193 

194 return cast(T, cached_value) 

195 

196 async with lock: 

197 # Check if another task filled the cache while we acquired the lock 

198 if (cached_value := cache_entry[key][0]) is initial_missing: 

199 self._misses += 1 

200 if self._maxsize is not None and self._currsize >= self._maxsize: 

201 cache_entry.popitem(last=False) 

202 else: 

203 self._currsize += 1 

204 

205 value = await self.__wrapped__(*args, **kwargs) 

206 expires_at = ( 

207 current_time() + self._ttl if self._ttl is not None else None 

208 ) 

209 cache_entry[key] = value, None, expires_at 

210 else: 

211 # Another task filled the cache while we were waiting for the lock 

212 self._hits += 1 

213 cache_entry.move_to_end(key) 

214 value = cast(T, cached_value) 

215 

216 return value 

217 

218 def __get__( 

219 self, instance: object, owner: type | None = None 

220 ) -> _LRUMethodWrapper[T]: 

221 wrapper = _LRUMethodWrapper(self, instance) 

222 update_wrapper(wrapper, self.__wrapped__) 

223 return wrapper 

224 

225 

226class _LRUCacheWrapper: 

227 def __init__( 

228 self, maxsize: int | None, typed: bool, always_checkpoint: bool, ttl: int | None 

229 ): 

230 self._maxsize = maxsize 

231 self._typed = typed 

232 self._always_checkpoint = always_checkpoint 

233 self._ttl = ttl 

234 

235 @overload 

236 def __call__( # type: ignore[overload-overlap] 

237 self, func: Callable[P, Coroutine[Any, Any, T]], / 

238 ) -> AsyncLRUCacheWrapper[P, T]: ... 

239 

240 @overload 

241 def __call__( 

242 self, func: Callable[..., T], / 

243 ) -> functools._lru_cache_wrapper[T]: ... 

244 

245 def __call__( 

246 self, f: Callable[P, Coroutine[Any, Any, T]] | Callable[..., T], / 

247 ) -> AsyncLRUCacheWrapper[P, T] | functools._lru_cache_wrapper[T]: 

248 if iscoroutinefunction(f): 

249 return AsyncLRUCacheWrapper( 

250 f, self._maxsize, self._typed, self._always_checkpoint, self._ttl 

251 ) 

252 

253 return functools.lru_cache(maxsize=self._maxsize, typed=self._typed)(f) # type: ignore[arg-type] 

254 

255 

256@overload 

257def cache( # type: ignore[overload-overlap] 

258 func: Callable[P, Coroutine[Any, Any, T]], / 

259) -> AsyncLRUCacheWrapper[P, T]: ... 

260 

261 

262@overload 

263def cache(func: Callable[..., T], /) -> functools._lru_cache_wrapper[T]: ... 

264 

265 

266def cache(func: Callable[..., Any] | Callable[P, Coroutine[Any, Any, Any]], /) -> Any: 

267 """ 

268 A convenient shortcut for :func:`lru_cache` with ``maxsize=None``. 

269 

270 This is the asynchronous equivalent to :func:`functools.cache`. 

271 

272 """ 

273 return lru_cache(maxsize=None)(func) 

274 

275 

276@overload 

277def lru_cache( 

278 *, 

279 maxsize: int | None = ..., 

280 typed: bool = ..., 

281 always_checkpoint: bool = ..., 

282 ttl: int | None = ..., 

283) -> _LRUCacheWrapper: ... 

284 

285 

286@overload 

287def lru_cache( # type: ignore[overload-overlap] 

288 func: Callable[P, Coroutine[Any, Any, T]], / 

289) -> AsyncLRUCacheWrapper[P, T]: ... 

290 

291 

292@overload 

293def lru_cache(func: Callable[..., T], /) -> functools._lru_cache_wrapper[T]: ... 

294 

295 

296def lru_cache( 

297 func: Callable[..., Coroutine[Any, Any, Any]] | Callable[..., Any] | None = None, 

298 /, 

299 *, 

300 maxsize: int | None = 128, 

301 typed: bool = False, 

302 always_checkpoint: bool = False, 

303 ttl: int | None = None, 

304) -> Any: 

305 """ 

306 An asynchronous version of :func:`functools.lru_cache`. 

307 

308 If a synchronous function is passed, the standard library 

309 :func:`functools.lru_cache` is applied instead. 

310 

311 :param always_checkpoint: if ``True``, every call to the cached function will be 

312 guaranteed to yield control to the event loop at least once 

313 :param ttl: time in seconds after which to invalidate cache entries 

314 

315 .. note:: Caches and locks are managed on a per-event loop basis. 

316 

317 """ 

318 if func is None: 

319 return _LRUCacheWrapper(maxsize, typed, always_checkpoint, ttl) 

320 

321 if not callable(func): 

322 raise TypeError("the first argument must be callable") 

323 

324 return _LRUCacheWrapper(maxsize, typed, always_checkpoint, ttl)(func) 

325 

326 

327@overload 

328async def reduce( 

329 function: Callable[[T, S], Awaitable[T]], 

330 iterable: Iterable[S] | AsyncIterable[S], 

331 /, 

332 initial: T, 

333) -> T: ... 

334 

335 

336@overload 

337async def reduce( 

338 function: Callable[[T, T], Awaitable[T]], 

339 iterable: Iterable[T] | AsyncIterable[T], 

340 /, 

341) -> T: ... 

342 

343 

344async def reduce( # type: ignore[misc] 

345 function: Callable[[T, T], Awaitable[T]] | Callable[[T, S], Awaitable[T]], 

346 iterable: Iterable[T] | Iterable[S] | AsyncIterable[T] | AsyncIterable[S], 

347 /, 

348 initial: T | _InitialMissingType = initial_missing, 

349) -> T: 

350 """ 

351 Asynchronous version of :func:`functools.reduce`. 

352 

353 :param function: a coroutine function that takes two arguments: the accumulated 

354 value and the next element from the iterable 

355 :param iterable: an iterable or async iterable 

356 :param initial: the initial value (if missing, the first element of the iterable is 

357 used as the initial value) 

358 

359 """ 

360 element: Any 

361 function_called = False 

362 if isinstance(iterable, AsyncIterable): 

363 async_it = iterable.__aiter__() 

364 if initial is initial_missing: 

365 try: 

366 value = cast(T, await async_it.__anext__()) 

367 except StopAsyncIteration: 

368 raise TypeError( 

369 "reduce() of empty sequence with no initial value" 

370 ) from None 

371 else: 

372 value = cast(T, initial) 

373 

374 async for element in async_it: 

375 value = await function(value, element) 

376 function_called = True 

377 elif isinstance(iterable, Iterable): 

378 it = iter(iterable) 

379 if initial is initial_missing: 

380 try: 

381 value = cast(T, next(it)) 

382 except StopIteration: 

383 raise TypeError( 

384 "reduce() of empty sequence with no initial value" 

385 ) from None 

386 else: 

387 value = cast(T, initial) 

388 

389 for element in it: 

390 value = await function(value, element) 

391 function_called = True 

392 else: 

393 raise TypeError("reduce() argument 2 must be an iterable or async iterable") 

394 

395 # Make sure there is at least one checkpoint, even if an empty iterable and an 

396 # initial value were given 

397 if not function_called: 

398 await checkpoint() 

399 

400 return value