from collections.abc import Callable from time import sleep from typing import ParamSpec, TypeVar import unittest P = ParamSpec("P") T = TypeVar("T") def retry( operation: Callable[P, T], *args: P.args, max_attempts: int = 3, initial_delay: float = 0.5, backoff_factor: float = 2.0, max_delay: float | None = None, retry_on: tuple[type[BaseException], ...] = (Exception,), sleeper: Callable[[float], None] = sleep, **kwargs: P.kwargs, ) -> T: """Call an operation, retrying failures with exponential backoff.""" if max_attempts < 1: raise ValueError("max_attempts must be at least 1") if initial_delay < 0: raise ValueError("initial_delay cannot be negative") if backoff_factor < 1: raise ValueError("backoff_factor must be at least 1") if max_delay is not None and max_delay < 0: raise ValueError("max_delay cannot be negative") delay = initial_delay for attempt in range(1, max_attempts + 1): try: return operation(*args, **kwargs) except retry_on: if attempt == max_attempts: raise # Wait before retrying, then increase the next delay. sleeper(min(delay, max_delay) if max_delay is not None else delay) delay *= backoff_factor raise RuntimeError("unreachable") class RetryTests(unittest.TestCase): def test_retries_until_success(self) -> None: attempts = 0 delays: list[float] = [] def flaky() -> str: nonlocal attempts attempts += 1 if attempts < 3: raise RuntimeError("temporary failure") return "success" result = retry( flaky, max_attempts=4, initial_delay=0.25, backoff_factor=2, sleeper=delays.append, ) self.assertEqual(result, "success") self.assertEqual(attempts, 3) self.assertEqual(delays, [0.25, 0.5]) def test_raises_after_attempts_are_exhausted(self) -> None: attempts = 0 def always_fails() -> None: nonlocal attempts attempts += 1 raise ValueError("still failing") with self.assertRaisesRegex(ValueError, "still failing"): retry(always_fails, max_attempts=3, sleeper=lambda _: None) self.assertEqual(attempts, 3) if __name__ == "__main__": unittest.main()