fix lint and some types

Author: jonathanedey

firebase_admin/_http_client.py

@@ -19,7 +19,7 @@
 
 from __future__ import annotations
 import logging
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, Generator, Optional, Tuple, Union
 import httpx
 import requests.adapters
 from requests.packages.urllib3.util import retry # pylint: disable=import-error
@@ -175,7 +175,10 @@ def __init__(self, credential: credentials.Credentials):
         self._refresh_status_codes = (401,)
 
     def apply_auth_headers(
-            self, request: httpx.Request, auth_request: google_auth_requests.Request):
+            self,
+            request: httpx.Request,
+            auth_request: google_auth_requests.Request
+        ) -> None:
         """A helper function to refreshes credentials if needed and mutates the request headers to
         contain access token and any other google auth headers."""
 
@@ -188,7 +191,7 @@ def apply_auth_headers(
         )
         logger.debug('Auth headers applied. Credential validity after: %s', self._credential.valid)
 
-    def auth_flow(self, request: httpx.Request):
+    def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]:
         _original_headers = request.headers.copy()
         _credential_refresh_attempt = 0
 

firebase_admin/_retry.py

@@ -23,8 +23,7 @@
 import random
 import re
 import time
-from types import CoroutineType
-from typing import Any, Callable, List, Optional, Tuple
+from typing import Any, Callable, List, Optional, Tuple, Coroutine
 import logging
 import asyncio
 import httpx
@@ -163,7 +162,7 @@ class HttpxRetryTransport(httpx.AsyncBaseTransport):
 
     DEFAULT_RETRY = HttpxRetry(max_retries=4, status_forcelist=[500, 503], backoff_factor=0.5)
 
-    def __init__(self, retry: HttpxRetry = DEFAULT_RETRY, **kwargs) -> None:
+    def __init__(self, retry: HttpxRetry = DEFAULT_RETRY, **kwargs: Any) -> None:
         self._retry = retry
 
         transport_kwargs = kwargs.copy()
@@ -180,7 +179,7 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
     async def _dispatch_with_retry(
             self,
             request: httpx.Request,
-            dispatch_method: Callable[[httpx.Request], CoroutineType[Any, Any, httpx.Response]]
+            dispatch_method: Callable[[httpx.Request], Coroutine[Any, Any, httpx.Response]]
     ) -> httpx.Response:
         """Sends a request with retry logic using a provided dispatch method."""
         # This request config is used across all requests that use this transport and therefore

firebase_admin/messaging.py

@@ -15,7 +15,7 @@
 """Firebase Cloud Messaging module."""
 
 from __future__ import annotations
-from typing import Callable, List, Optional
+from typing import Any, Callable, Dict, List, Optional, cast
 import concurrent.futures
 import json
 import warnings
@@ -37,7 +37,6 @@
     exceptions,
     App
 )
-from firebase_admin._retry import HttpxRetryTransport
 
 logger = logging.getLogger(__name__)
 
@@ -111,7 +110,7 @@
 def _get_messaging_service(app: Optional[App]) -> _MessagingService:
     return _utils.get_app_service(app, _MESSAGING_ATTRIBUTE, _MessagingService)
 
-def send(message, dry_run=False, app: Optional[App] = None):
+def send(message: Message, dry_run: bool = False, app: Optional[App] = None) -> str:
     """Sends the given message via Firebase Cloud Messaging (FCM).
 
     If the ``dry_run`` mode is enabled, the message will not be actually delivered to the
@@ -131,7 +130,11 @@ def send(message, dry_run=False, app: Optional[App] = None):
     """
     return _get_messaging_service(app).send(message, dry_run)
 
-def send_each(messages, dry_run=False, app=None):
+def send_each(
+        messages: List[Message],
+        dry_run: bool = False,
+        app: Optional[App] = None
+    ) -> BatchResponse:
     """Sends each message in the given list via Firebase Cloud Messaging.
 
     If the ``dry_run`` mode is enabled, the message will not be actually delivered to the
@@ -451,7 +454,7 @@ class _MessagingService:
         'UNREGISTERED': UnregisteredError,
     }
 
-    def __init__(self, app) -> None:
+    def __init__(self, app: App) -> None:
         project_id = app.project_id
         if not project_id:
             raise ValueError(
@@ -476,7 +479,7 @@ def encode_message(cls, message):
             raise ValueError('Message must be an instance of messaging.Message class.')
         return cls.JSON_ENCODER.default(message)
 
-    def send(self, message, dry_run=False):
+    def send(self, message: Message, dry_run: bool = False) -> str:
         """Sends the given message to FCM via the FCM v1 API."""
         data = self._message_data(message, dry_run)
         try:
@@ -489,9 +492,9 @@ def send(self, message, dry_run=False):
         except requests.exceptions.RequestException as error:
             raise self._handle_fcm_error(error)
         else:
-            return resp['name']
+            return cast(str, resp['name'])
 
-    def send_each(self, messages, dry_run=False):
+    def send_each(self, messages: List[Message], dry_run: bool = False) -> BatchResponse:
         """Sends the given messages to FCM via the FCM v1 API."""
         if not isinstance(messages, list):
             raise ValueError('messages must be a list of messaging.Message instances.')
@@ -683,7 +686,10 @@ def _build_fcm_error_requests(cls, error, message, error_dict):
 
     @classmethod
     def _build_fcm_error_httpx(
-                cls, error: httpx.HTTPError, message, error_dict
+            cls,
+            error: httpx.HTTPError,
+            message: str,
+            error_dict: Optional[Dict[str, Any]]
         ) -> Optional[exceptions.FirebaseError]:
         """Parses a httpx error response from the FCM API and creates a FCM-specific exception if
         appropriate."""
@@ -702,7 +708,11 @@ def _build_fcm_error_googleapiclient(cls, error, message, error_dict, http_respo
         return exc_type(message, cause=error, http_response=http_response) if exc_type else None
 
     @classmethod
-    def _build_fcm_error(cls, error_dict) -> Optional[Callable[..., exceptions.FirebaseError]]:
+    def _build_fcm_error(
+            cls,
+            error_dict: Optional[Dict[str, Any]]
+        ) -> Optional[Callable[..., exceptions.FirebaseError]]:
+        """Parses an error response to determine the appropriate FCM-specific error type."""
         if not error_dict:
             return None
         fcm_code = None