fix _ContextIterator and Context usage of SleepRetVal

kaioretry/context.py

@@ -93,7 +93,7 @@
 import logging
 import uuid
 
-from typing import cast, Awaitable, Any, TypeVar
+from typing import cast, Awaitable, Any, TypeVar, Generic
 from collections.abc import Callable, Generator, AsyncGenerator
 
 from .types import NonNegative, Number, UpdateDelayFunc
@@ -103,7 +103,7 @@
 SleepF = Callable[[Number], SleepRetVal]
 
 
-class _ContextIterator:
+class _ContextIterator(Generic[SleepRetVal]):
     """Single-usage helper class for Context objects."""
 
     # pylint: disable=too-few-public-methods
@@ -225,7 +225,7 @@ def __update_delay(self, delay: NonNegative) -> NonNegative:
         return delay
 
     def __make_iterator(
-            self, sleep: SleepF[Any]) -> Generator[SleepRetVal, None, None]:
+            self, sleep: SleepF[SleepRetVal]) -> Generator[SleepRetVal, None, None]:
         return iter(_ContextIterator(
             uuid.uuid4(), sleep, self.__tries, self.__delay,
             self.__update_delay, self.__logger))
@@ -236,8 +236,7 @@ def __iter__(self) -> Generator[None, None, None]:
 
     async def __aiter__(self) -> AsyncGenerator[None, None]:
         yield
-        for sleep in cast(Generator[Awaitable[None], None, None],
-                          self.__make_iterator(asyncio.sleep)):
+        for sleep in self.__make_iterator(asyncio.sleep):
             await sleep
             yield