Coverage for src/taipanstack/utils/cache.py: 100%

1""" 

2Intelligent Cache decorator. 

3 

4Provides in-memory caching that respects the Result monad and TTL, 

5ignoring caching for Err() results. 

6""" 

7 

8import asyncio 

9import functools 

10import inspect 

11import math 

12import time 

13from collections.abc import Awaitable, Callable 

14from typing import ParamSpec, Protocol, TypeAlias, TypeVar, cast, overload 

15 

16from taipanstack.core.result import Ok, Result 

17 

18P = ParamSpec("P") 

19T = TypeVar("T") 

20E = TypeVar("E", bound=Exception) 

21 

22CacheKey: TypeAlias = tuple[object, ...] 

23CacheValue: TypeAlias = tuple[float, object] 

24CacheDict: TypeAlias = dict[CacheKey, CacheValue] 

25 

26 

27def _check_cache( 

28 cache_key: CacheKey, 

29 cache: CacheDict, 

30 now: float, 

31) -> tuple[bool, object]: 

32 if cache_key in cache: 

33 expiry, value = cache[cache_key] 

34 if now < expiry: 

35 # Move to end to mark as recently used 

36 cache[cache_key] = cache.pop(cache_key) 

37 return True, value 

38 del cache[cache_key] 

39 return False, None 

40 

41 

42def _update_cache( 

43 cache_key: CacheKey, 

44 result: Result[T, E], 

45 cache: CacheDict, 

46 max_size: int, 

47 now: float, 

48 ttl: float, 

49) -> None: 

50 if isinstance(result, Ok): 

51 if len(cache) >= max_size: 

52 # Evict least recently used (first item) 

53 lru_key = next(iter(cache)) 

54 del cache[lru_key] 

55 cache[cache_key] = (now + ttl, result.ok_value) 

56 

57 

58def _hash_sequence(val: tuple[object, ...] | list[object]) -> tuple[object, ...]: 

59 return tuple(_make_hashable(item) for item in val) 

60 

61 

62def _hash_dict(val: dict[object, object]) -> tuple[object, ...]: 

63 items = [(_make_hashable(k), _make_hashable(v)) for k, v in val.items()] 

64 items.sort(key=lambda item: (type(item[0]).__name__, repr(item[0]))) # type: ignore[misc] 

65 return tuple(items) 

66 

67 

68def _hash_set(val: set[object]) -> frozenset[object]: 

69 return frozenset(_make_hashable(item) for item in val) 

70 

71 

72def _make_hashable(val: object) -> object: 

73 if isinstance(val, (tuple, list)): 

74 return _hash_sequence(val) 

75 elif isinstance(val, dict): 

76 return _hash_dict(val) 

77 elif isinstance(val, set): 

78 return _hash_set(val) 

79 else: 

80 hash(val) 

81 return val 

82 

83 

84def _get_cache_key( 

85 func_name: str, 

86 args: tuple[object, ...], 

87 kwargs: dict[str, object], 

88) -> CacheKey: 

89 hashable_args = tuple(_make_hashable(arg) for arg in args) 

90 hashable_kwargs = tuple( 

91 sorted((k, _make_hashable(v)) for k, v in kwargs.items()), 

92 ) 

93 return (func_name, hashable_args, hashable_kwargs) 

94 

95 

96def _is_valid_ttl_type(ttl: float) -> bool: 

97 return isinstance(ttl, (int, float)) and not isinstance(ttl, bool) 

98 

99 

100def _is_valid_ttl_value(ttl: float) -> bool: 

101 return math.isfinite(ttl) and ttl >= 0 

102 

103 

104def _validate_ttl(ttl: float) -> None: 

105 if not _is_valid_ttl_type(ttl) or not _is_valid_ttl_value(ttl): 

106 raise ValueError("ttl must be a finite non-negative number") 

107 

108 

109def _validate_max_size(max_size: int) -> None: 

110 if not isinstance(max_size, int) or isinstance(max_size, bool) or max_size <= 0: 

111 raise ValueError("max_size must be a positive integer") 

112 

113 

114def _validate_ttl_max_size(ttl: float, max_size: int) -> None: 

115 _validate_ttl(ttl) 

116 _validate_max_size(max_size) 

117 

118 

119def _get_or_create_lock( 

120 cache_key: CacheKey, 

121 locks: dict[CacheKey, asyncio.Lock], 

122 lock_waiters: dict[CacheKey, int], 

123) -> asyncio.Lock: 

124 if cache_key not in locks: 

125 locks[cache_key] = asyncio.Lock() 

126 lock_waiters[cache_key] = 0 

127 

128 lock_waiters[cache_key] += 1 

129 return locks[cache_key] 

130 

131 

132def _release_lock( 

133 cache_key: CacheKey, 

134 locks: dict[CacheKey, asyncio.Lock], 

135 lock_waiters: dict[CacheKey, int], 

136) -> None: 

137 lock_waiters[cache_key] -= 1 

138 if lock_waiters[cache_key] == 0: 

139 locks.pop(cache_key, None) 

140 lock_waiters.pop(cache_key, None) 

141 

142 

143def _execute_cached_async_wrapper( 

144 func: Callable[P, Awaitable[Result[T, E]]], 

145 cache: CacheDict, 

146 locks: dict[CacheKey, asyncio.Lock], 

147 lock_waiters: dict[CacheKey, int], 

148 max_size: int, 

149 ttl: float, 

150) -> Callable[P, Awaitable[Result[T, E]]]: 

151 @functools.wraps(func) 

152 async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Result[T, E]: 

153 cache_key = _get_cache_key( 

154 cast(str, getattr(func, "__name__", "unknown")), 

155 cast(tuple[object, ...], args), 

156 cast(dict[str, object], kwargs), 

157 ) 

158 

159 # Check cache before acquiring lock 

160 now = time.monotonic() 

161 hit, value = _check_cache(cache_key, cache, now) 

162 if hit: 

163 return Ok(cast(T, value)) 

164 

165 lock = _get_or_create_lock(cache_key, locks, lock_waiters) 

166 

167 try: 

168 async with lock: 

169 # Double-check cache after acquiring lock 

170 now = time.monotonic() 

171 hit, value = _check_cache(cache_key, cache, now) 

172 if hit: 

173 return Ok(cast(T, value)) 

174 

175 result = await func(*args, **kwargs) 

176 

177 _update_cache(cache_key, result, cache, max_size, now, ttl) 

178 return result 

179 finally: 

180 _release_lock(cache_key, locks, lock_waiters) 

181 

182 return async_wrapper # type: ignore[misc] 

183 

184 

185def _execute_cached_sync_wrapper( 

186 func: Callable[P, Result[T, E]], 

187 cache: CacheDict, 

188 max_size: int, 

189 ttl: float, 

190) -> Callable[P, Result[T, E]]: 

191 @functools.wraps(func) 

192 def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> Result[T, E]: 

193 cache_key = _get_cache_key( 

194 cast(str, getattr(func, "__name__", "unknown")), 

195 cast(tuple[object, ...], args), 

196 cast(dict[str, object], kwargs), 

197 ) 

198 now = time.monotonic() 

199 

200 hit, value = _check_cache(cache_key, cache, now) 

201 if hit: 

202 return Ok(cast(T, value)) 

203 

204 result = func(*args, **kwargs) 

205 

206 _update_cache(cache_key, result, cache, max_size, now, ttl) 

207 return result 

208 

209 return sync_wrapper 

210 

211 

212class CacheDecorator(Protocol): 

213 """Protocol for the cache decorator.""" 

214 

215 @overload 

216 def __call__( 

217 self, 

218 func: Callable[P, Result[T, E]], 

219 ) -> Callable[P, Result[T, E]]: ... 

220 

221 @overload 

222 def __call__( 

223 self, 

224 func: Callable[P, Awaitable[Result[T, E]]], 

225 ) -> Callable[P, Awaitable[Result[T, E]]]: ... 

226 

227 

228def cached(ttl: float, max_size: int = 1024) -> CacheDecorator: 

229 """Cache the Ok() results of a function for a given TTL. 

230 

231 Err() results are not cached. Supports both async and sync functions. 

232 Implements LRU (Least Recently Used) eviction when max_size is reached. 

233 

234 Args: 

235 ttl: Time to live in seconds. 

236 max_size: Maximum number of elements to store in the cache. 

237 

238 Returns: 

239 Decorator function. 

240 

241 """ 

242 _validate_ttl_max_size(ttl, max_size) 

243 

244 _cache: CacheDict = {} 

245 _locks: dict[CacheKey, asyncio.Lock] = {} 

246 _lock_waiters: dict[CacheKey, int] = {} 

247 

248 def decorator( 

249 func: Callable[P, Result[T, E]] | Callable[P, Awaitable[Result[T, E]]], 

250 ) -> Callable[P, Result[T, E]] | Callable[P, Awaitable[Result[T, E]]]: 

251 if inspect.iscoroutinefunction(func): 

252 func_coro = cast(Callable[P, Awaitable[Result[T, E]]], func) 

253 return _execute_cached_async_wrapper( 

254 func_coro, _cache, _locks, _lock_waiters, max_size, ttl 

255 ) 

256 

257 func_sync = cast(Callable[P, Result[T, E]], func) 

258 return _execute_cached_sync_wrapper(func_sync, _cache, max_size, ttl) 

259 

260 return cast(CacheDecorator, decorator)