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

1""" 

2Rate limiting utilities. 

3 

4Provides an in-memory token-bucket based rate limiting decorator 

5for both synchronous and asynchronous functions. The decorator 

6returns a ``Result`` type encapsulating the original return value 

7or a ``RateLimitError`` error. 

8""" 

9 

10import functools 

11import inspect 

12import math 

13import threading 

14import time 

15from collections.abc import Awaitable, Callable 

16from typing import ParamSpec, Protocol, TypeVar, cast, overload 

17 

18from taipanstack.core.result import Err, Ok, Result 

19 

20__all__ = ["RateLimitError", "RateLimiter", "rate_limit"] 

21 

22P = ParamSpec("P") 

23T = TypeVar("T") 

24 

25 

26class RateLimitError(Exception): 

27 """Exception raised when a rate limit is exceeded.""" 

28 

29 def __init__(self, message: str = "Rate limit exceeded") -> None: 

30 """Initialize the RateLimitError. 

31 

32 Args: 

33 message: The error message to display.Defaults to "Rate limit exceeded". 

34 

35 """ 

36 super().__init__(message) 

37 

38 

39class RateLimiter: 

40 """Token bucket rate limiter logic.""" 

41 

42 def __init__(self, max_calls: int, time_window: float) -> None: 

43 """Initialize the token bucket. 

44 

45 Args: 

46 max_calls: The maximum number of calls allowed in the time window. 

47 time_window: The time window in seconds. 

48 

49 """ 

50 if not math.isfinite(max_calls) or not math.isfinite(time_window): 

51 raise ValueError("max_calls and time_window must be finite numbers") 

52 if max_calls <= 0 or time_window <= 0: 

53 raise ValueError("max_calls and time_window must be > 0.0") 

54 self.capacity: float = float(max_calls) 

55 self.time_window: float = float(time_window) 

56 self.tokens: float = self.capacity 

57 self.last_update: float = time.monotonic() 

58 self._lock = threading.Lock() 

59 

60 def _is_valid_time_window(self) -> bool: 

61 """Check if time window is valid.""" 

62 if not isinstance(self.time_window, (int, float)): 

63 return False # type: ignore[unreachable] 

64 return math.isfinite(self.time_window) and self.time_window > 0.0 

65 

66 def _is_valid_capacity(self) -> bool: 

67 """Check if capacity is valid.""" 

68 if not isinstance(self.capacity, (int, float)): 

69 return False # type: ignore[unreachable] 

70 return math.isfinite(self.capacity) and self.capacity > 0.0 

71 

72 def _is_valid_bucket_state(self) -> bool: 

73 """Check if the bucket's time window and capacity are in a valid state.""" 

74 return self._is_valid_time_window() and self._is_valid_capacity() 

75 

76 def _calculate_new_tokens(self, elapsed: float) -> float | None: 

77 """Calculate new tokens based on elapsed time.""" 

78 new_tokens = elapsed * (self.capacity / self.time_window) 

79 return new_tokens if math.isfinite(new_tokens) else None 

80 

81 def _apply_new_tokens(self, new_tokens: float) -> bool: 

82 """Apply new tokens to the bucket.""" 

83 if not isinstance(self.tokens, (int, float)): 

84 self.tokens = self.capacity # type: ignore[unreachable] 

85 return False 

86 if not isinstance(new_tokens, (int, float)): 

87 return False # type: ignore[unreachable] 

88 self.tokens += new_tokens 

89 if not math.isfinite(self.tokens): 

90 # Reset to previous state or capacity if corrupted 

91 self.tokens = self.capacity 

92 return False 

93 self.tokens = min(self.tokens, self.capacity) 

94 return True 

95 

96 def _add_tokens(self, now: float) -> bool: 

97 """Calculate and add new tokens to the bucket based on elapsed time. 

98 

99 Args: 

100 now: Current monotonic time. 

101 

102 Returns: 

103 True if token update succeeds, False if state corruption is detected. 

104 

105 """ 

106 if not isinstance(self.last_update, (int, float)): 

107 return False # type: ignore[unreachable] 

108 raw_elapsed = now - self.last_update 

109 if not math.isfinite(raw_elapsed): 

110 return False 

111 elapsed = max(0.0, raw_elapsed) 

112 self.last_update = now 

113 

114 # Prevent state corruption or infinite elapsed time 

115 if not self._is_valid_bucket_state(): 

116 return False 

117 

118 new_tokens = self._calculate_new_tokens(elapsed) 

119 if new_tokens is None: 

120 return False 

121 

122 return self._apply_new_tokens(new_tokens) 

123 

124 def _try_consume(self, tokens: float) -> bool: 

125 """Attempt to consume the tokens from the bucket if available.""" 

126 if not isinstance(self.tokens, (int, float)): 

127 return False # type: ignore[unreachable] 

128 if not math.isfinite(self.tokens): 

129 # Reset to capacity if state is corrupted to inf/nan 

130 self.tokens = self.capacity 

131 return False 

132 if self.tokens >= tokens: 

133 self.tokens -= tokens 

134 return True 

135 return False 

136 

137 def _process_consumption(self, tokens: float) -> bool: 

138 """Process token consumption inside the lock.""" 

139 try: 

140 now = time.monotonic() 

141 except Exception: 

142 return False 

143 

144 # Prevent time corruption from poisoning the bucket state. 

145 # Only try to add tokens if time is finite. 

146 if not isinstance(now, (int, float)) or ( 

147 math.isfinite(now) and not self._add_tokens(now) 

148 ): 

149 return False 

150 

151 return self._try_consume(tokens) 

152 

153 def consume(self, tokens: float = 1.0) -> bool: 

154 """Try to consume tokens. 

155 

156 Args: 

157 tokens: Number of tokens to consume. Defaults to 1.0. 

158 

159 Returns: 

160 True if tokens were consumed (allow), False otherwise (limit exceeded). 

161 

162 """ 

163 if not isinstance(tokens, (int, float)) or not math.isfinite(tokens): 

164 return False 

165 if tokens <= 0: 

166 return True 

167 

168 try: 

169 with self._lock: 

170 return self._process_consumption(tokens) 

171 except Exception: 

172 return False 

173 

174 

175class RateLimitDecorator(Protocol): 

176 """Protocol for the rate limit decorator.""" 

177 

178 @overload 

179 def __call__( 

180 self, 

181 func: Callable[P, T], 

182 ) -> Callable[P, Result[T, RateLimitError]]: ... 

183 

184 @overload 

185 def __call__( 

186 self, 

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

188 ) -> Callable[P, Awaitable[Result[T, RateLimitError]]]: ... 

189 

190 

191def rate_limit( 

192 max_calls: int, 

193 time_window: float, 

194) -> RateLimitDecorator: 

195 """Decorate a function to apply rate limiting. 

196 

197 If the rate limit is exceeded, the wrapped function immediately returns 

198 an ``Err(RateLimitError)``. Uses an in-memory token bucket strategy. 

199 

200 Args: 

201 max_calls: Maximum function executions allowed in the defined window. 

202 time_window: Time window size in seconds. 

203 

204 Returns: 

205 Decorated function returning a ``Result[T, RateLimitError]``. 

206 

207 Example: 

208 >>> @rate_limit(max_calls=2, time_window=1.0) 

209 ... def fetch_data() -> str: 

210 ... return "data" 

211 >>> fetch_data() 

212 Ok('data') 

213 >>> fetch_data() 

214 Ok('data') 

215 >>> fetch_data() 

216 Err(RateLimitError('Rate limit exceeded')) 

217 

218 """ 

219 

220 def decorator( 

221 func: Callable[P, T] | Callable[P, Awaitable[T]], 

222 ) -> ( 

223 Callable[P, Result[T, RateLimitError]] 

224 | Callable[P, Awaitable[Result[T, RateLimitError]]] 

225 ): 

226 limiter = RateLimiter(max_calls, time_window) 

227 

228 if inspect.iscoroutinefunction(func): 

229 

230 @functools.wraps(func) 

231 async def async_wrapper( 

232 *args: P.args, 

233 **kwargs: P.kwargs, 

234 ) -> Result[T, RateLimitError]: 

235 try: 

236 if not limiter.consume(): 

237 return Err(RateLimitError()) 

238 except Exception: 

239 return Err(RateLimitError()) 

240 return Ok(await func(*args, **kwargs)) 

241 

242 return async_wrapper 

243 

244 @functools.wraps(func) 

245 def wrapper(*args: P.args, **kwargs: P.kwargs) -> Result[T, RateLimitError]: 

246 try: 

247 if not limiter.consume(): 

248 return Err(RateLimitError()) 

249 except Exception: 

250 return Err(RateLimitError()) 

251 func_sync = cast(Callable[P, T], func) 

252 return Ok(func_sync(*args, **kwargs)) 

253 

254 return wrapper 

255 

256 return cast(RateLimitDecorator, decorator)