Skip to content

Commit

Permalink
Add type hints to several functions
Browse files Browse the repository at this point in the history
  • Loading branch information
TomerYogev committed Jul 2, 2024
1 parent bce2043 commit f03bd49
Show file tree
Hide file tree
Showing 3 changed files with 34 additions and 26 deletions.
8 changes: 5 additions & 3 deletions aiogoogle/auth/creds.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
__all__ = ["ApiKey", "UserCreds", "ClientCreds", "ServiceAccountCreds"]


from typing import Mapping, Optional, Sequence
from ..utils import _dict


Expand Down Expand Up @@ -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
Expand Down
26 changes: 16 additions & 10 deletions aiogoogle/client.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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``
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
26 changes: 13 additions & 13 deletions aiogoogle/models.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit f03bd49

Please sign in to comment.