From e0a9365d612eb9dd6deeca3eebda7ddfe9e31b90 Mon Sep 17 00:00:00 2001 From: Josh Date: Thu, 8 Jul 2021 10:17:17 +1000 Subject: [PATCH] Type-hint backoff.py --- discord/backoff.py | 36 +++++++++++++++++++++++++++--------- 1 file changed, 27 insertions(+), 9 deletions(-) diff --git a/discord/backoff.py b/discord/backoff.py index f538facea..903ecf769 100644 --- a/discord/backoff.py +++ b/discord/backoff.py @@ -22,14 +22,20 @@ FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ +from __future__ import annotations + + import time import random +from typing import Callable, Generic, Literal, TypeVar, overload, Union + +T = TypeVar('T', bool, Literal[True], Literal[False]) __all__ = ( 'ExponentialBackoff', ) -class ExponentialBackoff: +class ExponentialBackoff(Generic[T]): """An implementation of the exponential backoff algorithm Provides a convenient interface to implement an exponential backoff @@ -51,21 +57,33 @@ class ExponentialBackoff: number in between may be returned. """ - def __init__(self, base=1, *, integral=False): - self._base = base + def __init__(self, base: int = 1, *, integral: T = False): + self._base: int = base - self._exp = 0 - self._max = 10 - self._reset_time = base * 2 ** 11 - self._last_invocation = time.monotonic() + self._exp: int = 0 + self._max: int = 10 + self._reset_time: int = base * 2 ** 11 + self._last_invocation: float = time.monotonic() # Use our own random instance to avoid messing with global one rand = random.Random() rand.seed() - self._randfunc = rand.randrange if integral else rand.uniform + self._randfunc: Callable[..., Union[int, float]] = rand.randrange if integral else rand.uniform # type: ignore + + @overload + def delay(self: ExponentialBackoff[Literal[False]]) -> float: + ... + + @overload + def delay(self: ExponentialBackoff[Literal[True]]) -> int: + ... + + @overload + def delay(self: ExponentialBackoff[bool]) -> Union[int, float]: + ... - def delay(self): + def delay(self) -> Union[int, float]: """Compute the next delay Returns the next delay to wait according to the exponential