from __future__ import annotations import time import unittest from collections.abc import Callable from typing import ParamSpec, TypeVar P = ParamSpec("P") T = TypeVar("T") def retry( function: Callable[P, T], *args: P.args, attempts: int = 3, initial_delay: float = 0.1, backoff_factor: float = 2.0, retry_on: tuple[type[Exception], ...] = (Exception,), sleep: Callable[[float], None] = time.sleep, **kwargs: P.kwargs, ) -> T: """Call a function until it succeeds or all attempts are exhausted.""" if attempts < 1: raise ValueError("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") delay = initial_delay for attempt in range(attempts): try: return function(*args, **kwargs) except retry_on: if attempt == attempts - 1: raise sleep(delay) delay *= backoff_factor raise RuntimeError("unreachable") class RetryTests(unittest.TestCase): def test_retries_with_exponential_backoff(self) -> None: calls = 0 delays: list[float] = [] def flaky() -> str: nonlocal calls calls += 1 if calls < 3: raise ConnectionError("temporary failure") return "ok" result = retry( flaky, attempts=4, initial_delay=0.5, backoff_factor=2, retry_on=(ConnectionError,), sleep=delays.append, ) self.assertEqual(result, "ok") self.assertEqual(calls, 3) self.assertEqual(delays, [0.5, 1.0]) def test_raises_after_final_attempt(self) -> None: def always_fails() -> None: raise ValueError("permanent failure") with self.assertRaisesRegex(ValueError, "permanent failure"): retry(always_fails, attempts=2, initial_delay=0, sleep=lambda _: None) if __name__ == "__main__": unittest.main()