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 time 

12from collections.abc import Awaitable, Callable 

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

14 

15from taipanstack.core.result import Ok, Result 

16 

17P = ParamSpec("P") 

18T = TypeVar("T") 

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

20 

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

22CacheValue: TypeAlias = tuple[float, object] 

23CacheDict: TypeAlias = dict[CacheKey, CacheValue] 

24 

25 

26def _check_cache( 

27 cache_key: CacheKey, 

28 cache: CacheDict, 

29 now: float, 

30) -> tuple[bool, object]: 

31 if cache_key in cache: 

32 expiry, value = cache[cache_key] 

33 if now < expiry: 

34 # Move to end to mark as recently used 

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

36 return True, value 

37 del cache[cache_key] 

38 return False, None 

39 

40 

41def _update_cache( 

42 cache_key: CacheKey, 

43 result: Result[T, E], 

44 cache: CacheDict, 

45 max_size: int, 

46 now: float, 

47 ttl: float, 

48) -> None: 

49 if isinstance(result, Ok): 

50 if len(cache) >= max_size: 

51 # Evict least recently used (first item) 

52 lru_key = next(iter(cache)) 

53 del cache[lru_key] 

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

55 

56 

57class CacheDecorator(Protocol): 

58 """Protocol for the cache decorator.""" 

59 

60 @overload 

61 def __call__( 

62 self, 

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

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

65 

66 @overload 

67 def __call__( 

68 self, 

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

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

71 

72 

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

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

75 

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

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

78 

79 Args: 

80 ttl: Time to live in seconds. 

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

82 

83 Returns: 

84 Decorator function. 

85 

86 """ 

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

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

89 

90 _cache: CacheDict = {} 

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

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

93 

94 def get_cache_key( 

95 func_name: str, 

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

97 kwargs: dict[str, object], 

98 ) -> CacheKey: 

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

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

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

102 elif isinstance(val, dict): 

103 return tuple(sorted((k, _make_hashable(v)) for k, v in val.items())) 

104 elif isinstance(val, set): 

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

106 else: 

107 hash(val) 

108 return val 

109 

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

111 hashable_kwargs = tuple( 

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

113 ) 

114 return (func_name, hashable_args, hashable_kwargs) 

115 

116 def decorator( 

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

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

119 if inspect.iscoroutinefunction(func): 

120 

121 @functools.wraps(func) 

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

123 cache_key = get_cache_key( 

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

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

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

127 ) 

128 

129 # Check cache before acquiring lock 

130 now = time.monotonic() 

131 hit, value = _check_cache(cache_key, _cache, now) 

132 if hit: 

133 return Ok(cast(T, value)) 

134 

135 if cache_key not in _locks: 

136 _locks[cache_key] = asyncio.Lock() 

137 _lock_waiters[cache_key] = 0 

138 

139 _lock_waiters[cache_key] += 1 

140 lock = _locks[cache_key] 

141 

142 try: 

143 async with lock: 

144 # Double-check cache after acquiring lock 

145 now = time.monotonic() 

146 hit, value = _check_cache(cache_key, _cache, now) 

147 if hit: 

148 return Ok(cast(T, value)) 

149 

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

151 result = await func_coro(*args, **kwargs) 

152 

153 _update_cache(cache_key, result, _cache, max_size, now, ttl) 

154 return result 

155 finally: 

156 _lock_waiters[cache_key] -= 1 

157 if _lock_waiters[cache_key] == 0: 

158 _locks.pop(cache_key, None) 

159 _lock_waiters.pop(cache_key, None) 

160 

161 return async_wrapper 

162 

163 @functools.wraps(func) 

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

165 cache_key = get_cache_key( 

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

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

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

169 ) 

170 now = time.monotonic() 

171 

172 hit, value = _check_cache(cache_key, _cache, now) 

173 if hit: 

174 return Ok(cast(T, value)) 

175 

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

177 result = func_sync(*args, **kwargs) 

178 

179 _update_cache(cache_key, result, _cache, max_size, now, ttl) 

180 return result 

181 

182 return sync_wrapper 

183 

184 return cast(CacheDecorator, decorator)