diff --git a/backoff/_wait_gen.py b/backoff/_wait_gen.py index 3fa5d54..cc9c885 100644 --- a/backoff/_wait_gen.py +++ b/backoff/_wait_gen.py @@ -8,7 +8,7 @@ def expo( base: float = 2, factor: float = 1, max_value: Optional[float] = None -) -> Generator[Optional[float], Any, None]: +) -> Generator[float, Any, None]: """Generator for exponential decay. @@ -54,7 +54,7 @@ def fibo(max_value: Optional[int] = None) -> Generator[int, None, None]: def constant( interval: Union[int, Iterable[float]] = 1 -) -> Generator[Optional[float], None, None]: +) -> Generator[float, None, None]: """Generator for constant intervals. Args: @@ -75,7 +75,7 @@ def constant( def runtime( *, value: Callable[[Any], float] -) -> Generator[Optional[float], None, None]: +) -> Generator[float, None, None]: """Generator that is based on parsing the return value or thrown exception of the decorated method diff --git a/tests/test_typing.py b/tests/test_typing.py new file mode 100644 index 0000000..7f53459 --- /dev/null +++ b/tests/test_typing.py @@ -0,0 +1,34 @@ +import backoff + + +# No pyunit tests are defined here yet, but the following decorator calls will +# be analyzed by mypy which would have caught a bug the last release. + +@backoff.on_exception( + backoff.expo, + ValueError, + jitter=None, + max_tries=3, +) +def foo(): + raise ValueError() + + +@backoff.on_exception( + backoff.constant, + ValueError, + interval=1, + max_tries=3 +) +def bar(): + raise ValueError() + + +@backoff.on_predicate( + backoff.runtime, + predicate=lambda r: r.status_code == 429, + value=lambda r: int(r.headers.get("Retry-After")), + jitter=None, +) +def baz(): + pass