fixed mypy issues

Author: daniel-sanche

google/api_core/retry_streaming.py

@@ -24,6 +24,7 @@
     Generator,
     TypeVar,
     Any,
+    Union,
     cast,
 )
 
@@ -38,13 +39,6 @@
 T = TypeVar("T")
 
 
-class _TerminalException(Exception):
-    """
-    Exception to bypasses retry logic and raises __cause__ immediately.
-    """
-    pass
-
-
 def _build_timeout_error(
     exc_list: List[Exception], is_timeout: bool, timeout_val: float
 ) -> Tuple[Exception, Optional[Exception]]:
@@ -74,15 +68,13 @@ def _build_timeout_error(
 
 
 def retry_target_generator(
-    target: Callable[[], Iterable[T]],
+    target: Callable[[], Union[Iterable[T], Generator[T, Any, None]]],
     predicate: Callable[[Exception], bool],
     sleep_generator: Iterable[float],
     timeout: Optional[float] = None,
     on_error: Optional[Callable[[Exception], None]] = None,
     exception_factory: Optional[
-        Callable[
-            [List[Exception], bool, float], Tuple[Exception, Optional[Exception]]
-        ]
+        Callable[[List[Exception], bool, float], Tuple[Exception, Optional[Exception]]]
     ] = None,
     **kwargs,
 ) -> Generator[T, Any, None]:
@@ -171,17 +163,19 @@ def on_error(e):
         except Exception as exc:
             error_list.append(exc)
             if not predicate(exc):
-                exc, source_exc = exc_factory(exc_list=error_list, is_timeout=False)
-                raise exc from source_exc
+                final_exc, source_exc = exc_factory(
+                    exc_list=error_list, is_timeout=False
+                )
+                raise final_exc from source_exc
             if on_error is not None:
                 on_error(exc)
         finally:
             if subgenerator is not None and getattr(subgenerator, "close", None):
-                subgenerator.close()
+                cast(Generator, subgenerator).close()
 
         if deadline is not None and time.monotonic() + sleep > deadline:
-            exc, source_exc = exc_factory(exc_list=error_list, is_timeout=True)
-            raise exc from source_exc
+            final_exc, source_exc = exc_factory(exc_list=error_list, is_timeout=True)
+            raise final_exc from source_exc
         _LOGGER.debug(
             "Retrying due to {}, sleeping {:.1f}s ...".format(error_list[-1], sleep)
         )

google/api_core/retry_streaming_async.py

@@ -42,13 +42,6 @@
 T = TypeVar("T")
 
 
-class _TerminalException(Exception):
-    """
-    Exception to bypasses retry logic and raises __cause__ immediately.
-    """
-    pass
-
-
 async def retry_target_generator(
     target: Union[
         Callable[[], AsyncIterable[T]],
@@ -133,12 +126,12 @@ def on_error(e):
             filter_retry_wrapped = retryable_with_filter(target)
             ```
     """
-    subgenerator = None
-
+    subgenerator : Optional[AsyncIterator[T]] = None
     timeout = kwargs.get("deadline", timeout)
-
     deadline: Optional[float] = time.monotonic() + timeout if timeout else None
+    # keep track of retryable exceptions we encounter to pass in to exception_factory
     error_list: List[Exception] = []
+    # override exception_factory to build a more complex exception
     exc_factory = partial(
         exception_factory or _build_timeout_error, timeout_val=timeout
     )
@@ -147,12 +140,13 @@ def on_error(e):
         # Start a new retry loop
         try:
             # generator may be raw iterator, or wrapped in an awaitable
-            subgenerator = target()
+            gen_instance: Union[AsyncIterable[T], Awaitable[AsyncIterable[T]]] = target()
             try:
-                subgenerator = await subgenerator
+                gen_instance = await gen_instance # type: ignore
             except TypeError:
                 # was not awaitable
                 pass
+            subgenerator = cast(AsyncIterable[T], gen_instance).__aiter__()
 
             # if target is a generator, we will advance it using asend
             # otherwise, we will use anext
@@ -162,7 +156,7 @@ def on_error(e):
             while True:
                 ## Read from Subgenerator
                 if supports_send:
-                    next_value = await subgenerator.asend(sent_in)
+                    next_value = await subgenerator.asend(sent_in) # type: ignore
                 else:
                     next_value = await subgenerator.__anext__()
                 ## Yield from Wrapper to caller
@@ -173,20 +167,18 @@ def on_error(e):
                 except GeneratorExit:
                     # if wrapper received `aclose`, pass to subgenerator and close
                     if bool(getattr(subgenerator, "aclose", None)):
-                        await subgenerator.aclose()
+                        await cast(AsyncGenerator[T, None], subgenerator).aclose()
                     else:
                         raise
                     return
                 except:  # noqa: E722
                     # bare except catches any exception passed to `athrow`
                     # delegate error handling to subgenerator
                     if getattr(subgenerator, "athrow", None):
-                        await subgenerator.athrow(*sys.exc_info())
+                        await cast(AsyncGenerator[T, None], subgenerator).athrow(*sys.exc_info())
                     else:
                         raise
             return
-        except _TerminalException as exc:
-            raise exc.__cause__ from exc.__cause__.__cause__
         except StopAsyncIteration:
             # if generator exhausted, return
             return
@@ -201,12 +193,12 @@ def on_error(e):
                 on_error(exc)
         finally:
             if subgenerator is not None and getattr(subgenerator, "aclose", None):
-                await subgenerator.aclose()
+                await cast(AsyncGenerator[T, None], subgenerator).aclose()
 
         # sleep and adjust timeout budget
         if deadline is not None and time.monotonic() + sleep > deadline:
-            exc, source_exc = exc_factory(exc_list=error_list, is_timeout=True)
-            raise exc from source_exc
+            final_exc, source_exc = exc_factory(exc_list=error_list, is_timeout=True)
+            raise final_exc from source_exc
         _LOGGER.debug(
             "Retrying due to {}, sleeping {:.1f}s ...".format(error_list[-1], sleep)
         )

tests/asyncio/test_retry_async.py

@@ -705,8 +705,10 @@ async def __anext__(self):
             return CustomIterable(n)
 
         if awaitale_wrapped:
+
             async def wrapper(n):
                 return iterable_fn(n)
+
             decorated = retry_(wrapper)
         else:
             decorated = retry_(iterable_fn)
@@ -718,7 +720,6 @@ async def wrapper(n):
         await retryable.asend("test2") == 2
         await retryable.asend("test3") == 3
 
-
     @pytest.mark.parametrize("awaitale_wrapped", [True, False])
     @mock.patch("asyncio.sleep", autospec=True)
     @pytest.mark.asyncio
@@ -744,9 +745,12 @@ async def __anext__(self):
                     return self.i - 1
 
             return CustomIterable(n)
+
         if awaitale_wrapped:
+
             async def wrapper(n):
                 return iterable_fn(n)
+
             decorated = retry_(wrapper)
         else:
             decorated = retry_(iterable_fn)
@@ -763,7 +767,6 @@ async def wrapper(n):
         with pytest.raises(StopAsyncIteration):
             await new_retryable.__anext__()
 
-
     @pytest.mark.parametrize("awaitale_wrapped", [True, False])
     @mock.patch("asyncio.sleep", autospec=True)
     @pytest.mark.asyncio
@@ -791,9 +794,12 @@ async def __anext__(self):
                     return self.i - 1
 
             return CustomIterable(n)
+
         if awaitale_wrapped:
+
             async def wrapper(n):
                 return iterable_fn(n)
+
             decorated = retry_(wrapper)
         else:
             decorated = retry_(iterable_fn)

tests/unit/test_retry.py

@@ -668,7 +668,6 @@ def iterable_fn(n):
         with pytest.raises(AttributeError):
             generator.send("test")
 
-
     @mock.patch("time.sleep", autospec=True)
     def test___call___with_iterable_close(self, sleep):
         """