Coverage for src/taipanstack/resilience/resilience.py: 100%

1""" 

2Resilience decorators. 

3 

4Provides tools for graceful fallback and timeouts using the Result monad. 

5""" 

6 

7import asyncio 

8import functools 

9import inspect 

10import math 

11import threading 

12from collections.abc import Awaitable, Callable 

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

14 

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

16 

17P = ParamSpec("P") 

18T = TypeVar("T") 

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

20 

21ResultFunc: TypeAlias = Callable[P, Result[T, E]] 

22AsyncResultFunc: TypeAlias = Callable[P, Awaitable[Result[T, E]]] 

23 

24 

25class FallbackDecorator(Protocol): 

26 """Protocol for the fallback decorator.""" 

27 

28 @overload 

29 def __call__(self, func: ResultFunc[P, T, E]) -> ResultFunc[P, T, E]: ... 

30 

31 @overload 

32 def __call__(self, func: AsyncResultFunc[P, T, E]) -> AsyncResultFunc[P, T, E]: ... 

33 

34 

35def _handle_fallback_exception( 

36 e: Exception, 

37 exceptions: tuple[type[Exception], ...], 

38 fallback_value: T, 

39) -> Result[T, E] | None: 

40 try: 

41 if isinstance(e, exceptions): 

42 return Ok(fallback_value) 

43 except TypeError: 

44 pass 

45 return None 

46 

47 

48def _process_fallback_result( 

49 result: Result[T, E], fallback_value: T 

50) -> Result[T, E] | None: 

51 if isinstance(result, Err): 

52 return Ok(fallback_value) 

53 if isinstance(result, Ok): 

54 return result 

55 return None # type: ignore[unreachable] 

56 

57 

58async def _run_fallback_async( 

59 func_coro: AsyncResultFunc[P, T, E], 

60 fallback_value: T, 

61 exceptions: tuple[type[Exception], ...], 

62 *args: P.args, 

63 **kwargs: P.kwargs, 

64) -> Result[T, E]: 

65 try: 

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

67 processed_res: Result[T, E] | None = _process_fallback_result( 

68 result, fallback_value 

69 ) 

70 if processed_res is not None: 

71 return processed_res 

72 except Exception as e: 

73 fallback_res: Result[T, E] | None = _handle_fallback_exception( 

74 e, exceptions, fallback_value 

75 ) 

76 if fallback_res is not None: 

77 return fallback_res 

78 raise 

79 return Err(cast(E, RuntimeError("Unreachable"))) 

80 

81 

82def _execute_fallback_async_wrapper( 

83 func_coro: AsyncResultFunc[P, T, E], 

84 fallback_value: T, 

85 exceptions: tuple[type[Exception], ...], 

86) -> AsyncResultFunc[P, T, E]: 

87 @functools.wraps(func_coro) 

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

89 return await _run_fallback_async( 

90 func_coro, fallback_value, exceptions, *args, **kwargs 

91 ) 

92 

93 return async_wrapper # type: ignore[misc] 

94 

95 

96def _run_fallback_sync( 

97 func_sync: ResultFunc[P, T, E], 

98 fallback_value: T, 

99 exceptions: tuple[type[Exception], ...], 

100 *args: P.args, 

101 **kwargs: P.kwargs, 

102) -> Result[T, E]: 

103 try: 

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

105 processed_res: Result[T, E] | None = _process_fallback_result( 

106 result, fallback_value 

107 ) 

108 if processed_res is not None: 

109 return processed_res 

110 except Exception as e: 

111 fallback_res: Result[T, E] | None = _handle_fallback_exception( 

112 e, exceptions, fallback_value 

113 ) 

114 if fallback_res is not None: 

115 return fallback_res 

116 raise 

117 return Err(cast(E, RuntimeError("Unreachable"))) 

118 

119 

120def _execute_fallback_sync_wrapper( 

121 func_sync: ResultFunc[P, T, E], 

122 fallback_value: T, 

123 exceptions: tuple[type[Exception], ...], 

124) -> ResultFunc[P, T, E]: 

125 @functools.wraps(func_sync) 

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

127 return _run_fallback_sync( 

128 func_sync, fallback_value, exceptions, *args, **kwargs 

129 ) 

130 

131 return sync_wrapper 

132 

133 

134def fallback( 

135 fallback_value: T, 

136 exceptions: tuple[type[Exception], ...] = (Exception,), 

137) -> FallbackDecorator: 

138 """Provide a fallback value on failures. 

139 

140 If the wrapped function returns an Err() or raises a specified exception, 

141 the fallback value is returned wrapped in an Ok(). 

142 

143 Args: 

144 fallback_value: The value to return on failure. 

145 exceptions: Exceptions to catch. 

146 

147 Returns: 

148 Decorator function. 

149 

150 """ 

151 

152 def decorator( 

153 func: ResultFunc[P, T, E] | AsyncResultFunc[P, T, E], 

154 ) -> ResultFunc[P, T, E] | AsyncResultFunc[P, T, E]: 

155 if inspect.iscoroutinefunction(func): 

156 func_coro = cast(AsyncResultFunc[P, T, E], func) 

157 return _execute_fallback_async_wrapper( 

158 func_coro, fallback_value, exceptions 

159 ) 

160 

161 func_sync = cast(ResultFunc[P, T, E], func) 

162 return _execute_fallback_sync_wrapper(func_sync, fallback_value, exceptions) 

163 

164 return cast(FallbackDecorator, decorator) 

165 

166 

167class TimeoutDecorator(Protocol): 

168 """Protocol for the timeout decorator.""" 

169 

170 @overload 

171 def __call__( 

172 self, 

173 func: ResultFunc[P, T, E], 

174 ) -> Callable[P, Result[T, TimeoutError | E]]: ... 

175 

176 @overload 

177 def __call__( 

178 self, 

179 func: AsyncResultFunc[P, T, E], 

180 ) -> Callable[P, Awaitable[Result[T, TimeoutError | E]]]: ... 

181 

182 

183def _handle_timeout_exception( 

184 e: BaseException, context: str 

185) -> Result[T, TimeoutError | E]: 

186 if isinstance(e, (SystemExit, KeyboardInterrupt, GeneratorExit)): 

187 raise e 

188 if isinstance(e, MemoryError): 

189 return Err(cast(E, RuntimeError(f"Memory exhaustion: {e!s}"))) 

190 if isinstance(e, (OSError, OverflowError)): 

191 return Err(cast(E, RuntimeError(f"Resource exhaustion: {e!s}"))) 

192 return Err(cast(E, RuntimeError(f"{context} exhaustion: {e!s}"))) 

193 

194 

195def _is_valid_timeout_value(seconds: float) -> bool: 

196 """Check if timeout is a valid finite non-negative number.""" 

197 if type(seconds) not in (int, float) or isinstance(seconds, bool): 

198 return False 

199 return bool(math.isfinite(seconds) and seconds >= 0) 

200 

201 

202def _validate_timeout(seconds: float) -> Result[None, E] | None: 

203 if not _is_valid_timeout_value(seconds): 

204 return Err( 

205 cast( 

206 E, 

207 ValueError("Timeout must be a finite non-negative number"), 

208 ), 

209 ) 

210 return None 

211 

212 

213async def _run_async_with_timeout( 

214 func_coro: Callable[P, Awaitable[Result[T, TimeoutError | E]]], 

215 seconds: float, 

216 *args: P.args, 

217 **kwargs: P.kwargs, 

218) -> Result[T, TimeoutError | E]: 

219 try: 

220 return await asyncio.wait_for( 

221 func_coro(*args, **kwargs), 

222 timeout=seconds, 

223 ) 

224 except TimeoutError: 

225 return Err(TimeoutError(f"Execution timed out after {seconds} seconds.")) 

226 except asyncio.CancelledError: 

227 raise 

228 except BaseException as e: 

229 return _handle_timeout_exception(e, "Task") 

230 

231 

232def _execute_timeout_async_wrapper( 

233 func_coro: Callable[P, Awaitable[Result[T, TimeoutError | E]]], 

234 seconds: float, 

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

236 @functools.wraps(func_coro) 

237 async def async_wrapper( 

238 *args: P.args, 

239 **kwargs: P.kwargs, 

240 ) -> Result[T, TimeoutError | E]: 

241 val_err: Result[None, E] | None = _validate_timeout(seconds) 

242 if val_err is not None: 

243 return cast(Result[T, TimeoutError | E], val_err) 

244 return await _run_async_with_timeout(func_coro, seconds, *args, **kwargs) 

245 

246 return async_wrapper # type: ignore[misc] 

247 

248 

249def _process_sync_timeout_result( 

250 seconds: float, 

251 thread: threading.Thread, 

252 result: list[Result[T, TimeoutError | E]], 

253 exception: list[BaseException], 

254) -> Result[T, TimeoutError | E]: 

255 if thread.is_alive(): 

256 return Err(TimeoutError(f"Execution timed out after {seconds} seconds.")) 

257 

258 if exception: 

259 raise exception[0] 

260 

261 if not result: 

262 err_msg = "Thread died abruptly without a result or exception" 

263 return Err(cast(E, RuntimeError(err_msg))) 

264 

265 return result[0] 

266 

267 

268def _run_sync_with_timeout( 

269 func_sync: Callable[P, Result[T, TimeoutError | E]], 

270 seconds: float, 

271 *args: P.args, 

272 **kwargs: P.kwargs, 

273) -> Result[T, TimeoutError | E]: 

274 result: list[Result[T, TimeoutError | E]] = [] 

275 exception: list[BaseException] = [] 

276 

277 def worker() -> None: 

278 try: 

279 result.append(func_sync(*args, **kwargs)) 

280 except BaseException as e: 

281 exception.append(e) 

282 

283 try: 

284 thread = threading.Thread(target=worker, daemon=True) 

285 thread.start() 

286 thread.join(timeout=seconds) 

287 except BaseException as e: 

288 return _handle_timeout_exception(e, "Thread") 

289 

290 return _process_sync_timeout_result(seconds, thread, result, exception) 

291 

292 

293def _execute_timeout_sync_wrapper( 

294 func_sync: Callable[P, Result[T, TimeoutError | E]], 

295 seconds: float, 

296) -> Callable[P, Result[T, TimeoutError | E]]: 

297 @functools.wraps(func_sync) 

298 def sync_wrapper( 

299 *args: P.args, 

300 **kwargs: P.kwargs, 

301 ) -> Result[T, TimeoutError | E]: 

302 val_err: Result[None, E] | None = _validate_timeout(seconds) 

303 if val_err is not None: 

304 return cast(Result[T, TimeoutError | E], val_err) 

305 return _run_sync_with_timeout(func_sync, seconds, *args, **kwargs) 

306 

307 return sync_wrapper 

308 

309 

310def timeout(seconds: float) -> TimeoutDecorator: 

311 """Enforce a maximum execution time. 

312 

313 If the execution time exceeds the specified limit, returns Err(TimeoutError). 

314 

315 Args: 

316 seconds: Maximum allowed execution time in seconds. 

317 

318 Returns: 

319 Decorator function. 

320 

321 """ 

322 

323 def decorator( 

324 func: ResultFunc[P, T, E] | AsyncResultFunc[P, T, E], 

325 ) -> ( 

326 Callable[P, Result[T, TimeoutError | E]] 

327 | Callable[P, Awaitable[Result[T, TimeoutError | E]]] 

328 ): 

329 if inspect.iscoroutinefunction(func): 

330 func_coro = cast(Callable[P, Awaitable[Result[T, TimeoutError | E]]], func) 

331 return _execute_timeout_async_wrapper(func_coro, seconds) 

332 

333 func_sync = cast(Callable[P, Result[T, TimeoutError | E]], func) 

334 return _execute_timeout_sync_wrapper(func_sync, seconds) 

335 

336 return cast(TimeoutDecorator, decorator)