support timeout as aiohttpClientTimeout and total_attempts (max retries)

Author: yinghsienwu

google/auth/aio/transport/aiohttp.py

@@ -17,7 +17,8 @@
 
 import asyncio
 import logging
-from typing import AsyncGenerator, Mapping, Optional
+import typing
+from typing import Any, AsyncGenerator, Mapping, Optional, Union
 
 try:
     import aiohttp  # type: ignore
@@ -26,6 +27,15 @@
         "The aiohttp library is not installed from please install the aiohttp package to use the aiohttp transport."
     ) from caught_exc
 
+if typing.TYPE_CHECKING:
+    from aiohttp import ClientTimeout
+else:
+    ClientTimeout: typing.Type = Any
+    try:
+        from aiohttp import ClientTimeout
+    except ImportError:
+        ClientTimeout = None
+
 from google.auth import _helpers
 from google.auth import exceptions
 from google.auth.aio import _helpers as _helpers_async
@@ -123,7 +133,7 @@ async def __call__(
         method: str = "GET",
         body: Optional[bytes] = None,
         headers: Optional[Mapping[str, str]] = None,
-        timeout: float = transport._DEFAULT_TIMEOUT_SECONDS,
+        timeout: Union[float, ClientTimeout] = transport._DEFAULT_TIMEOUT_SECONDS,
         **kwargs,
     ) -> transport.Response:
         """
@@ -158,7 +168,10 @@ async def __call__(
             if not self._session:
                 self._session = aiohttp.ClientSession()
 
-            client_timeout = aiohttp.ClientTimeout(total=timeout)
+            if isinstance(timeout, aiohttp.ClientTimeout):
+                client_timeout = timeout
+            else:
+              client_timeout = aiohttp.ClientTimeout(total=timeout)
             _helpers.request_log(_LOGGER, method, url, body, headers)
             response = await self._session.request(
                 method,

google/auth/aio/transport/sessions.py

@@ -16,7 +16,8 @@
 from contextlib import asynccontextmanager
 import functools
 import time
-from typing import Mapping, Optional
+import typing
+from typing import Any, Mapping, Optional, Union
 
 from google.auth import _exponential_backoff, exceptions
 from google.auth.aio import transport
@@ -30,7 +31,16 @@
 except ImportError:  # pragma: NO COVER
     AIOHTTP_INSTALLED = False
 
-
+if typing.TYPE_CHECKING:
+    from aiohttp import ClientTimeout
+else:
+    ClientTimeout: typing.Type = Any
+    try:
+        from aiohttp import ClientTimeout
+    except ImportError:
+        ClientTimeout = None
+        
+    
 @asynccontextmanager
 async def timeout_guard(timeout):
     """
@@ -137,7 +147,8 @@ async def request(
         data: Optional[bytes] = None,
         headers: Optional[Mapping[str, str]] = None,
         max_allowed_time: float = transport._DEFAULT_TIMEOUT_SECONDS,
-        timeout: float = transport._DEFAULT_TIMEOUT_SECONDS,
+        timeout: Union[float, ClientTimeout] = transport._DEFAULT_TIMEOUT_SECONDS,
+        total_attempts: Optional[int] = transport.DEFAULT_MAX_RETRY_ATTEMPTS,
         **kwargs,
     ) -> transport.Response:
         """
@@ -146,14 +157,16 @@ async def request(
                 url (str): The URI to be requested.
                 data (Optional[bytes]): The payload or body in HTTP request.
                 headers (Optional[Mapping[str, str]]): Request headers.
-                timeout (float):
+                timeout (float, aiohttp.ClientTimeout):
                 The amount of time in seconds to wait for the server response
                 with each individual request.
                 max_allowed_time (float):
                 If the method runs longer than this, a ``Timeout`` exception is
                 automatically raised. Unlike the ``timeout`` parameter, this
                 value applies to the total method execution time, even if
                 multiple requests are made under the hood.
+                total_attempts (int):
+                The total number of retry attempts.
 
                 Mind that it is not guaranteed that the timeout error is raised
                 at ``max_allowed_time``. It might take longer, for example, if
@@ -172,7 +185,7 @@ async def request(
         """
 
         retries = _exponential_backoff.AsyncExponentialBackoff(
-            total_attempts=transport.DEFAULT_MAX_RETRY_ATTEMPTS
+            total_attempts=total_attempts,
         )
         async with timeout_guard(max_allowed_time) as with_timeout:
             await with_timeout(
@@ -198,7 +211,7 @@ async def get(
         data: Optional[bytes] = None,
         headers: Optional[Mapping[str, str]] = None,
         max_allowed_time: float = transport._DEFAULT_TIMEOUT_SECONDS,
-        timeout: float = transport._DEFAULT_TIMEOUT_SECONDS,
+        timeout: Union[float, ClientTimeout] = transport._DEFAULT_TIMEOUT_SECONDS,
         **kwargs,
     ) -> transport.Response:
         return await self.request(
@@ -212,7 +225,7 @@ async def post(
         data: Optional[bytes] = None,
         headers: Optional[Mapping[str, str]] = None,
         max_allowed_time: float = transport._DEFAULT_TIMEOUT_SECONDS,
-        timeout: float = transport._DEFAULT_TIMEOUT_SECONDS,
+        timeout: Union[float, ClientTimeout] = transport._DEFAULT_TIMEOUT_SECONDS,
         **kwargs,
     ) -> transport.Response:
         return await self.request(
@@ -226,7 +239,7 @@ async def put(
         data: Optional[bytes] = None,
         headers: Optional[Mapping[str, str]] = None,
         max_allowed_time: float = transport._DEFAULT_TIMEOUT_SECONDS,
-        timeout: float = transport._DEFAULT_TIMEOUT_SECONDS,
+        timeout: Union[float, ClientTimeout] = transport._DEFAULT_TIMEOUT_SECONDS,
         **kwargs,
     ) -> transport.Response:
         return await self.request(
@@ -240,7 +253,7 @@ async def patch(
         data: Optional[bytes] = None,
         headers: Optional[Mapping[str, str]] = None,
         max_allowed_time: float = transport._DEFAULT_TIMEOUT_SECONDS,
-        timeout: float = transport._DEFAULT_TIMEOUT_SECONDS,
+        timeout: Union[float, ClientTimeout] = transport._DEFAULT_TIMEOUT_SECONDS,
         **kwargs,
     ) -> transport.Response:
         return await self.request(