From 6cb1a4ecaf6905fe46069a0676e25c5bd5b6883a Mon Sep 17 00:00:00 2001
From: Bob Green <rgreen@aquent.com>
Date: Wed, 5 Oct 2022 15:11:17 -0400
Subject: [PATCH] Fix type hint for wait generator

I picked this up in one of the type hinting improvement PRs but it
doesn't seem to work and isn't necessary.

This also adds a test file which will be picked up by mypy when
running make check which would have caught the error.

Fixes https://github.com/litl/backoff/issues/177
---
 backoff/_wait_gen.py |  6 +++---
 tests/test_typing.py | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_typing.py

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