Merge pull request #146 from TomerYogev/add-type-annotations

Add type hints to several functions

Author: omarryhan

aiogoogle/auth/creds.py

@@ -1,6 +1,7 @@
 __all__ = ["ApiKey", "UserCreds", "ClientCreds", "ServiceAccountCreds"]
 
 
+from typing import Mapping, Optional, Sequence
 from ..utils import _dict
 
 
@@ -273,9 +274,10 @@ class ServiceAccountCreds(_dict):
             }
     """
     def __init__(
-        self, type=None, project_id=None, private_key_id=None, private_key=None,
-        client_email=None, client_id=None, auth_uri=None, token_uri=None, auth_provider_x509_cert_url=None,
-        client_x509_cert_url=None, subject=None, scopes=None, additional_claims=None, universe_domain='googleapis.com'
+        self, type: Optional[str] = None, project_id: Optional[str] = None, private_key_id: Optional[str] = None, private_key: Optional[str] = None,
+        client_email: Optional[str] = None, client_id: Optional[str] = None, auth_uri: Optional[str] = None, token_uri: Optional[str] = None, auth_provider_x509_cert_url: Optional[str] = None,
+        client_x509_cert_url: Optional[str] = None, subject: Optional[str] = None, scopes: Optional[Sequence[str]] = None, additional_claims: Optional[Mapping[str, str]] = None,
+        universe_domain: str = 'googleapis.com'
     ):
         self.type = type
         self.project_id = project_id

aiogoogle/client.py

@@ -1,13 +1,19 @@
+from __future__ import annotations
+
 __all__ = ["Aiogoogle"]
 
 from contextvars import ContextVar
-from typing import Optional
+from typing import TYPE_CHECKING, Any, Optional, Type
 
 from .resource import GoogleAPI
 from .auth.managers import Oauth2Manager, ApiKeyManager, OpenIdConnectManager, ServiceAccountManager
 from .sessions.aiohttp_session import AiohttpSession
 from .data import DISCOVERY_SERVICE_V1_DISCOVERY_DOC
 
+if TYPE_CHECKING:
+    from .auth.creds import ApiKey, ClientCreds, ServiceAccountCreds, UserCreds
+    from .models import Request, Response
+    from .sessions.abc import AbstractSession
 
 # Discovery doc reference https://developers.google.com/discovery/v1/reference/apis
 
@@ -52,11 +58,11 @@ class Aiogoogle:
 
     def __init__(
         self,
-        session_factory=AiohttpSession,
-        api_key=None,
-        user_creds=None,
-        client_creds=None,
-        service_account_creds=None,
+        session_factory: Type[AbstractSession] = AiohttpSession,
+        api_key: Optional[ApiKey] = None,
+        user_creds: Optional[UserCreds] = None,
+        client_creds: Optional[ClientCreds] = None,
+        service_account_creds: Optional[ServiceAccountCreds] = None,
     ):
 
         self.session_factory = session_factory
@@ -147,7 +153,7 @@ async def list_api(self, name, preferred=None, fields=None):
         )
         return await self.as_anon(request)
 
-    async def discover(self, api_name, api_version=None, validate=False, *, disco_doc_ver: Optional[int] = None):
+    async def discover(self, api_name: str, api_version: Optional[str] = None, validate: bool = False, *, disco_doc_ver: Optional[int] = None) -> GoogleAPI:
         """
         Donwloads a discovery document from Google's Discovery Service V1 and sets it a ``aiogoogle.resource.GoogleAPI``
 
@@ -262,7 +268,7 @@ async def as_user(self, *requests, timeout=None, full_res=False, user_creds=None
         )
 
     async def as_service_account(
-            self, *requests, timeout=None, full_res=False, service_account_creds=None, raise_for_status=True):
+            self, *requests: Request, timeout: Optional[int] = None, full_res: bool = False, service_account_creds: ServiceAccountCreds = None, raise_for_status: bool = True) -> Response:
         """
         Sends requests on behalf of ``self.user_creds`` (OAuth2)
 
@@ -408,15 +414,15 @@ async def send(self, *args, **kwargs):
             session = self._set_session()
         return await session.send(*args, **kwargs)
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> Aiogoogle:
         session = self._get_session()
         if session is None:
             session = self._set_session()
             await session.__aenter__()
             return self
         raise RuntimeError("Nesting context managers using the same Aiogoogle object is not allowed.")
 
-    async def __aexit__(self, *args):
+    async def __aexit__(self, *args: Any) -> None:
         session = self._get_session()
         await session.__aexit__(*args)
         # Had to add this because there's no use of keeping a closed session

aiogoogle/models.py

@@ -1,5 +1,5 @@
 from urllib.parse import urlparse, urlunparse, urlencode, parse_qs
-from typing import AsyncIterable
+from typing import Any, AsyncIterable, Callable, Optional
 import pprint
 
 from .excs import HTTPError, AuthError, ValidationError
@@ -179,18 +179,18 @@ class Request:
 
     def __init__(
         self,
-        method=None,
-        url=None,
-        batch_url=None,
-        headers=None,
-        json=None,
-        data=None,
-        media_upload=None,
-        media_download=None,
-        timeout=None,
-        callback=None,
-        _verify_ssl=True,
-        upload_file_content_type=None,
+        method: Optional[str] = None,
+        url: Optional[str] = None,
+        batch_url: Optional[str] = None,
+        headers: Optional[dict] = None,
+        json: Optional[dict] = None,
+        data: Any = None,
+        media_upload: Optional[MediaUpload] = None,
+        media_download: Optional[MediaDownload] = None,
+        timeout: Optional[int] = None,
+        callback: Optional[Callable] = None,
+        _verify_ssl: bool = True,
+        upload_file_content_type: Optional[str] = None,
     ):
         self.method = method
         self.url = url