Coverage for src/lexigram/graphql/dataloader/cache.py: 78%

1"""DataLoaderProtocol cache implementations. 

2 

3This module provides cache implementations for DataLoaderProtocol, 

4including in-memory caching with TTL support. 

5""" 

6 

7from __future__ import annotations 

8 

9from abc import ABC, abstractmethod 

10from collections.abc import Hashable 

11from dataclasses import dataclass 

12import time 

13from typing import Generic, TypeVar 

14 

15from lexigram.logging import get_logger 

16 

17logger = get_logger(__name__) 

18 

19K = TypeVar("K", bound=Hashable) 

20V = TypeVar("V") 

21 

22 

23class LoaderCache(ABC, Generic[K, V]): 

24 """Abstract base class for DataLoaderProtocol caches. 

25 

26 Implement this interface to provide custom caching 

27 behavior for DataLoaders. 

28 """ 

29 

30 @abstractmethod 

31 def get(self, key: K) -> V | None: 

32 """Get a cached value. 

33 

34 Args: 

35 key: Cache key. 

36 

37 Returns: 

38 Cached value or None. 

39 """ 

40 ... 

41 

42 @abstractmethod 

43 def set(self, key: K, value: V) -> None: 

44 """Set a cached value. 

45 

46 Args: 

47 key: Cache key. 

48 value: Value to cache. 

49 """ 

50 ... 

51 

52 @abstractmethod 

53 def has(self, key: K) -> bool: 

54 """Check if key exists in cache. 

55 

56 Args: 

57 key: Cache key. 

58 

59 Returns: 

60 True if key exists. 

61 """ 

62 ... 

63 

64 @abstractmethod 

65 def delete(self, key: K) -> None: 

66 """Delete a cached value. 

67 

68 Args: 

69 key: Cache key. 

70 """ 

71 ... 

72 

73 @abstractmethod 

74 def clear(self) -> None: 

75 """Clear all cached values.""" 

76 ... 

77 

78 

79@dataclass 

80class LoaderCacheEntry(Generic[V]): 

81 """Cache entry with expiration support. 

82 

83 Attributes: 

84 value: Cached value. 

85 expires_at: Expiration timestamp (or 0 for no expiry). 

86 """ 

87 

88 value: V 

89 expires_at: float = 0 

90 

91 def is_expired(self) -> bool: 

92 """Check if entry is expired.""" 

93 if self.expires_at <= 0: 

94 return False 

95 return time.time() > self.expires_at 

96 

97 

98class InMemoryCache(LoaderCache[K, V]): 

99 """In-memory cache implementation. 

100 

101 Simple in-memory cache with optional TTL support. 

102 Suitable for single-request caching in DataLoaders. 

103 

104 Example: 

105 ```python 

106 cache = InMemoryCache[str, User](ttl_seconds=60) 

107 

108 cache.set("user:1", user) 

109 user = cache.get("user:1") 

110 ``` 

111 """ 

112 

113 def __init__( 

114 self, 

115 ttl_seconds: float = 0, 

116 max_size: int = 0, 

117 ) -> None: 

118 """Initialize the cache. 

119 

120 Args: 

121 ttl_seconds: Time-to-live in seconds (0 for no TTL). 

122 max_size: Maximum cache size (0 for unlimited). 

123 """ 

124 self._ttl_seconds = ttl_seconds 

125 self._max_size = max_size 

126 self._cache: dict[K, LoaderCacheEntry[V]] = {} 

127 

128 @property 

129 def size(self) -> int: 

130 """Get current cache size.""" 

131 return len(self._cache) 

132 

133 def get(self, key: K) -> V | None: 

134 """Get a cached value. 

135 

136 Args: 

137 key: Cache key. 

138 

139 Returns: 

140 Cached value or None if not found or expired. 

141 """ 

142 entry = self._cache.get(key) 

143 

144 if entry is None: 

145 return None 

146 

147 if entry.is_expired(): 

148 del self._cache[key] 

149 return None 

150 

151 return entry.value 

152 

153 def set(self, key: K, value: V) -> None: 

154 """Set a cached value. 

155 

156 Args: 

157 key: Cache key. 

158 value: Value to cache. 

159 """ 

160 # Check max size 

161 if self._max_size > 0 and len(self._cache) >= self._max_size and self._cache: 

162 # Remove oldest entry (FIFO eviction) 

163 oldest_key = next(iter(self._cache)) 

164 del self._cache[oldest_key] 

165 

166 # Calculate expiration 

167 expires_at = 0.0 

168 if self._ttl_seconds > 0: 

169 expires_at = time.time() + self._ttl_seconds 

170 

171 self._cache[key] = LoaderCacheEntry(value=value, expires_at=expires_at) 

172 

173 def has(self, key: K) -> bool: 

174 """Check if key exists and is not expired. 

175 

176 Args: 

177 key: Cache key. 

178 

179 Returns: 

180 True if key exists and is valid. 

181 """ 

182 entry = self._cache.get(key) 

183 

184 if entry is None: 

185 return False 

186 

187 if entry.is_expired(): 

188 del self._cache[key] 

189 return False 

190 

191 return True 

192 

193 def delete(self, key: K) -> None: 

194 """Delete a cached value. 

195 

196 Args: 

197 key: Cache key. 

198 """ 

199 self._cache.pop(key, None) 

200 

201 def clear(self) -> None: 

202 """Clear all cached values.""" 

203 self._cache.clear() 

204 

205 def cleanup_expired(self) -> int: 

206 """Remove expired entries. 

207 

208 Returns: 

209 Number of entries removed. 

210 """ 

211 expired_keys = [ 

212 kv[0] for kv in filter(lambda kv: kv[1].is_expired(), self._cache.items()) 

213 ] 

214 

215 for key in expired_keys: 

216 del self._cache[key] 

217 

218 return len(expired_keys) 

219 

220 

221class NoOpCache(LoaderCache[K, V]): 

222 """No-operation cache (disables caching). 

223 

224 Use this when you want to disable DataLoaderProtocol caching 

225 while keeping the batching behavior. 

226 """ 

227 

228 def get(self, key: K) -> V | None: 

229 """Always returns None.""" 

230 return None 

231 

232 def set(self, key: K, value: V) -> None: 

233 """Does nothing.""" 

234 

235 def has(self, key: K) -> bool: 

236 """Always returns False.""" 

237 return False 

238 

239 def delete(self, key: K) -> None: 

240 """Does nothing.""" 

241 

242 def clear(self) -> None: 

243 """Does nothing.""" 

244 

245 

246__all__ = ["InMemoryCache", "LoaderCache", "LoaderCacheEntry", "NoOpCache"]