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 fallback( 

36 fallback_value: T, 

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

38) -> FallbackDecorator: 

39 """Provide a fallback value on failures. 

40 

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

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

43 

44 Args: 

45 fallback_value: The value to return on failure. 

46 exceptions: Exceptions to catch. 

47 

48 Returns: 

49 Decorator function. 

50 

51 """ 

52 

53 def decorator( 

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

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

56 if inspect.iscoroutinefunction(func): 

57 

58 @functools.wraps(func) 

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

60 try: 

61 # func is a coroutine function here 

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

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

64 if isinstance(result, Err): 

65 return Ok(fallback_value) 

66 elif isinstance(result, Ok): 

67 return result 

68 except Exception as e: 

69 try: 

70 if isinstance(e, exceptions): 

71 return Ok(fallback_value) 

72 except TypeError: 

73 pass 

74 raise 

75 return Err(cast(E, RuntimeError("Unreachable"))) # type: ignore[unreachable] 

76 

77 return async_wrapper 

78 

79 @functools.wraps(func) 

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

81 try: 

82 # func is a normal function here 

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

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

85 if isinstance(result, Err): 

86 return Ok(fallback_value) 

87 elif isinstance(result, Ok): 

88 return result 

89 except Exception as e: 

90 try: 

91 if isinstance(e, exceptions): 

92 return Ok(fallback_value) 

93 except TypeError: 

94 pass 

95 raise 

96 return Err(cast(E, RuntimeError("Unreachable"))) # type: ignore[unreachable] 

97 

98 return sync_wrapper 

99 

100 return cast(FallbackDecorator, decorator) 

101 

102 

103class TimeoutDecorator(Protocol): 

104 """Protocol for the timeout decorator.""" 

105 

106 @overload 

107 def __call__( 

108 self, 

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

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

111 

112 @overload 

113 def __call__( 

114 self, 

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

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

117 

118 

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

120 """Enforce a maximum execution time. 

121 

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

123 

124 Args: 

125 seconds: Maximum allowed execution time in seconds. 

126 

127 Returns: 

128 Decorator function. 

129 

130 """ 

131 

132 def decorator( 

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

134 ) -> ( 

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

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

137 ): 

138 if inspect.iscoroutinefunction(func): 

139 

140 @functools.wraps(func) 

141 async def async_wrapper( 

142 *args: P.args, 

143 **kwargs: P.kwargs, 

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

145 if ( 

146 not isinstance(seconds, (int, float)) 

147 or not math.isfinite(seconds) 

148 or seconds < 0 

149 ): 

150 return Err( 

151 cast( 

152 E, 

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

154 ), 

155 ) 

156 

157 try: 

158 func_coro = cast( 

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

160 func, 

161 ) 

162 return await asyncio.wait_for( 

163 func_coro(*args, **kwargs), 

164 timeout=seconds, 

165 ) 

166 except TimeoutError: 

167 return Err( 

168 TimeoutError(f"Execution timed out after {seconds} seconds."), 

169 ) 

170 except RuntimeError as e: 

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

172 except MemoryError as e: 

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

174 except (OSError, OverflowError) as e: 

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

176 

177 return async_wrapper 

178 

179 @functools.wraps(func) 

180 def sync_wrapper( 

181 *args: P.args, 

182 **kwargs: P.kwargs, 

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

184 if ( 

185 not isinstance(seconds, (int, float)) 

186 or not math.isfinite(seconds) 

187 or seconds < 0 

188 ): 

189 return Err( 

190 cast( 

191 E, 

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

193 ), 

194 ) 

195 

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

197 exception: list[BaseException] = [] 

198 

199 def worker() -> None: 

200 try: 

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

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

203 except BaseException as e: 

204 exception.append(e) 

205 

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

207 try: 

208 thread.start() 

209 thread.join(timeout=seconds) 

210 except RuntimeError as e: 

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

212 except (OSError, OverflowError) as e: 

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

214 except MemoryError as e: 

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

216 

217 if thread.is_alive(): 

218 return Err( 

219 TimeoutError(f"Execution timed out after {seconds} seconds."), 

220 ) 

221 

222 if exception: 

223 raise exception[0] 

224 

225 return result[0] 

226 

227 return sync_wrapper 

228 

229 return cast(TimeoutDecorator, decorator)