improved type hinting

Author: daniel-sanche

google/api_core/retry_streaming.py

@@ -14,29 +14,29 @@
 
 """Helpers for retries for streaming APIs."""
 
-from typing import Callable, Optional, Iterable
+from typing import Callable, Optional, Iterable, Iterator, Generator, TypeVar, Any, cast
 
 import datetime
 import logging
 import time
 
-from collections.abc import Generator
-
 from google.api_core import datetime_helpers
 from google.api_core import exceptions
 
 _LOGGER = logging.getLogger(__name__)
 
+T = TypeVar("T")
+
 
-class RetryableGenerator(Generator):
+class RetryableGenerator(Generator[T, Any, None]):
     """
     Helper class for retrying Iterator and Generator-based
     streaming APIs.
     """
 
     def __init__(
         self,
-        target: Callable[[], Iterable],
+        target: Callable[[], Iterable[T]],
         predicate: Callable[[Exception], bool],
         sleep_generator: Iterable[float],
         timeout: Optional[float] = None,
@@ -52,13 +52,13 @@ def __init__(
                 It should return True to retry or False otherwise.
             sleep_generator: An infinite iterator that determines
                 how long to sleep between retries.
-            timeout: How long to keep retrying the target.
+            timeout: How long to keep retrying the target, in seconds.
             on_error: A function to call while processing a
                 retryable exception.  Any error raised by this function will *not*
                 be caught.
         """
         self.target_fn = target
-        self.active_target: Iterable = self.target_fn()
+        self.active_target: Iterator = self.target_fn().__iter__()
         self.predicate = predicate
         self.sleep_generator = iter(sleep_generator)
         self.on_error = on_error
@@ -70,13 +70,13 @@ def __init__(
         else:
             self.deadline = None
 
-    def __iter__(self):
+    def __iter__(self) -> Generator[T, Any, None]:
         """
         Implement the iterator protocol.
         """
         return self
 
-    def _handle_exception(self, exc):
+    def _handle_exception(self, exc) -> None:
         """
         When an exception is raised while iterating over the active_target,
         check if it is retryable. If so, create a new active_target and
@@ -108,9 +108,12 @@ def _handle_exception(self, exc):
             time.sleep(next_sleep)
             self.active_target = self.target_fn()
 
-    def __next__(self):
+    def __next__(self) -> T:
         """
         Implement the iterator protocol.
+
+        Returns:
+          - the next value of the active_target iterator
         """
         try:
             return next(self.active_target)
@@ -119,7 +122,7 @@ def __next__(self):
         # if retryable exception was handled, try again with new active_target
         return self.__next__()
 
-    def close(self):
+    def close(self) -> None:
         """
         Close the active_target if supported. (e.g. target is a generator)
 
@@ -133,22 +136,25 @@ def close(self):
                 "close() not implemented for {}".format(self.active_target)
             )
 
-    def send(self, value):
+    def send(self, *args, **kwargs) -> T:
         """
         Call send on the active_target if supported. (e.g. target is a generator)
 
         If an exception is raised, a retry may be attempted before returning
         a result.
 
+        Args:
+          - *args: arguments to pass to the wrapped generator's send method
+          - **kwargs: keyword arguments to pass to the wrapped generator's send method
         Returns:
-          - the result of calling send() on the active_target
-
+          - the next value of the active_target iterator after calling send
         Raises:
           - AttributeError if the active_target does not have a send() method
         """
         if getattr(self.active_target, "send", None):
+            casted_target = cast(Generator, self.active_target)
             try:
-                return self.active_target.send(value)
+                return casted_target.send(*args, **kwargs)
             except Exception as exc:
                 self._handle_exception(exc)
             # if exception was retryable, use new target for return value
@@ -158,21 +164,25 @@ def send(self, value):
                 "send() not implemented for {}".format(self.active_target)
             )
 
-    def throw(self, typ, val=None, tb=None):
+    def throw(self, *args, **kwargs) -> T:
         """
         Call throw on the active_target if supported. (e.g. target is a generator)
 
         If an exception is raised, a retry may be attempted before returning
         a result.
 
+        Args:
+          - *args: arguments to pass to the wrapped generator's throw method
+          - **kwargs: keyword arguments to pass to the wrapped generator's throw method
         Returns:
-          - the result of calling throw() on the active_target
+          - the next vale of the active_target iterator after calling throw
         Raises:
           - AttributeError if the active_target does not have a throw() method
         """
         if getattr(self.active_target, "throw", None):
+            casted_target = cast(Generator, self.active_target)
             try:
-                return self.active_target.throw(typ, val, tb)
+                return casted_target.throw(*args, **kwargs)
             except Exception as exc:
                 self._handle_exception(exc)
             # if retryable exception was handled, return next from new active_target

google/api_core/retry_streaming_async.py

@@ -14,21 +14,35 @@
 
 """Helpers for retries for async streaming APIs."""
 
-from typing import Callable, Optional, Iterable, AsyncIterable, Awaitable, Union
+from typing import (
+    cast,
+    Callable,
+    Optional,
+    Iterable,
+    AsyncIterator,
+    AsyncIterable,
+    Awaitable,
+    Union,
+    Any,
+    TypeVar,
+    AsyncGenerator,
+)
 
 import asyncio
 import inspect
 import logging
+import datetime
 
-from collections.abc import AsyncGenerator
 
 from google.api_core import datetime_helpers
 from google.api_core import exceptions
 
 _LOGGER = logging.getLogger(__name__)
 
+T = TypeVar("T")
 
-class AsyncRetryableGenerator(AsyncGenerator):
+
+class AsyncRetryableGenerator(AsyncGenerator[T, None]):
     """
     Helper class for retrying AsyncIterator and AsyncGenerator-based
     streaming APIs.
@@ -37,7 +51,8 @@ class AsyncRetryableGenerator(AsyncGenerator):
     def __init__(
         self,
         target: Union[
-            Callable[[], AsyncIterable], Callable[[], Awaitable[AsyncIterable]]
+            Callable[[], AsyncIterable[T]],
+            Callable[[], Awaitable[AsyncIterable[T]]],
         ],
         predicate: Callable[[Exception], bool],
         sleep_generator: Iterable[float],
@@ -61,27 +76,32 @@ def __init__(
         """
         self.target_fn = target
         # active target must be populated in an async context
-        self.active_target: Optional[AsyncIterable] = None
+        self.active_target: Optional[AsyncIterator[T]] = None
         self.predicate = predicate
         self.sleep_generator = iter(sleep_generator)
         self.on_error = on_error
         self.timeout = timeout
         self.remaining_timeout_budget = timeout if timeout else None
 
-    async def _ensure_active_target(self):
+    async def _ensure_active_target(self) -> AsyncIterator[T]:
         """
         Ensure that the active target is populated and ready to be iterated over.
+
+        Returns:
+          - The active_target iterable
         """
         if not self.active_target:
-            self.active_target = self.target_fn()
-            if inspect.iscoroutine(self.active_target):
-                self.active_target = await self.active_target
+            new_iterable = self.target_fn()
+            if isinstance(new_iterable, Awaitable):
+                new_iterable = await new_iterable
+            self.active_target = new_iterable.__aiter__()
+        return self.active_target
 
-    def __aiter__(self):
+    def __aiter__(self) -> AsyncIterator[T]:
         """Implement the async iterator protocol."""
         return self
 
-    async def _handle_exception(self, exc):
+    async def _handle_exception(self, exc) -> None:
         """
         When an exception is raised while iterating over the active_target,
         check if it is retryable. If so, create a new active_target and
@@ -114,7 +134,7 @@ async def _handle_exception(self, exc):
             self.active_target = None
             await self._ensure_active_target()
 
-    def _subtract_time_from_budget(self, start_timestamp):
+    def _subtract_time_from_budget(self, start_timestamp: datetime.datetime) -> None:
         """
         Subtract the time elapsed since start_timestamp from the remaining
         timeout budget.
@@ -128,13 +148,15 @@ def _subtract_time_from_budget(self, start_timestamp):
                 datetime_helpers.utcnow() - start_timestamp
             ).total_seconds()
 
-    async def _iteration_helper(self, iteration_routine: Awaitable):
+    async def _iteration_helper(self, iteration_routine: Awaitable) -> T:
         """
         Helper function for sharing logic between __anext__ and asend.
 
         Args:
           - iteration_routine: The coroutine to await to get the next value
                 from the iterator (e.g. __anext__ or asend)
+        Returns:
+          - The next value from the active_target iterator.
         """
         # check for expired timeouts before attempting to iterate
         if (
@@ -164,16 +186,19 @@ async def _iteration_helper(self, iteration_routine: Awaitable):
         # if retryable exception was handled, find the next value to return
         return await self.__anext__()
 
-    async def __anext__(self):
+    async def __anext__(self) -> T:
         """
         Implement the async iterator protocol.
+
+        Returns:
+          - The next value from the active_target iterator.
         """
-        await self._ensure_active_target()
+        iterable = await self._ensure_active_target()
         return await self._iteration_helper(
-            self.active_target.__anext__(),
+            iterable.__anext__(),
         )
 
-    async def aclose(self):
+    async def aclose(self) -> None:
         """
         Close the active_target if supported. (e.g. target is an async generator)
 
@@ -182,48 +207,57 @@ async def aclose(self):
         """
         await self._ensure_active_target()
         if getattr(self.active_target, "aclose", None):
-            return await self.active_target.aclose()
+            casted_target = cast(AsyncGenerator[T, None], self.active_target)
+            return await casted_target.aclose()
         else:
             raise AttributeError(
                 "aclose() not implemented for {}".format(self.active_target)
             )
 
-    async def asend(self, value):
+    async def asend(self, *args, **kwargs) -> T:
         """
         Call asend on the active_target if supported. (e.g. target is an async generator)
 
         If an exception is raised, a retry may be attempted before returning
         a result.
 
-        Returns:
-          - the result of calling asend() on the active_target
 
+        Args:
+          - *args: arguments to pass to the wrapped generator's asend method
+          - **kwargs: keyword arguments to pass to the wrapped generator's asend method
+        Returns:
+          - the next value of the active_target iterator after calling asend
         Raises:
           - AttributeError if the active_target does not have a asend() method
         """
         await self._ensure_active_target()
         if getattr(self.active_target, "asend", None):
-            return await self._iteration_helper(self.active_target.asend(value))
+            casted_target = cast(AsyncGenerator[T, None], self.active_target)
+            return await self._iteration_helper(casted_target.asend(*args, **kwargs))
         else:
             raise AttributeError(
                 "asend() not implemented for {}".format(self.active_target)
             )
 
-    async def athrow(self, typ, val=None, tb=None):
+    async def athrow(self, *args, **kwargs) -> T:
         """
         Call athrow on the active_target if supported. (e.g. target is an async generator)
 
         If an exception is raised, a retry may be attempted before returning
 
+        Args:
+          - *args: arguments to pass to the wrapped generator's athrow method
+          - **kwargs: keyword arguments to pass to the wrapped generator's athrow method
         Returns:
-          - the result of calling athrow() on the active_target
+          - the next value of the active_target iterator after calling athrow
         Raises:
           - AttributeError if the active_target does not have a athrow() method
         """
         await self._ensure_active_target()
         if getattr(self.active_target, "athrow", None):
+            casted_target = cast(AsyncGenerator[T, None], self.active_target)
             try:
-                return await self.active_target.athrow(typ, val, tb)
+                return await casted_target.athrow(*args, **kwargs)
             except Exception as exc:
                 await self._handle_exception(exc)
             # if retryable exception was handled, return next from new active_target