From 48d9a399e04e05a2e2bd3846ea54aadea1f89464 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fran=C3=A7ois=20Voron?= Date: Sun, 7 Feb 2021 09:58:52 +0100 Subject: [PATCH] Fix type errors --- fastapi_users/db/ormar.py | 23 +++++++++++------------ fastapi_users/db/tortoise.py | 5 +++-- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/fastapi_users/db/ormar.py b/fastapi_users/db/ormar.py index 6b675c4c..878e258a 100644 --- a/fastapi_users/db/ormar.py +++ b/fastapi_users/db/ormar.py @@ -1,5 +1,4 @@ -import datetime -from typing import Any, List, Optional, Type +from typing import cast, Any, List, Optional, Type import ormar from ormar.exceptions import NoMatch @@ -74,8 +73,8 @@ class OrmarUserDatabase(BaseUserDatabase[UD]): model = await self.model(**user.dict(exclude={"oauth_accounts"})).save() if oauth_accounts and self.oauth_account_model: await self._create_oauth_models(model=model, oauth_accounts=oauth_accounts) - user = await self._get_user(id=user.id) - return user + user_db = await self._get_user(id=user.id) + return cast(UD, user_db) async def update(self, user: UD) -> UD: oauth_accounts = getattr(user, "oauth_accounts", []) @@ -84,27 +83,27 @@ class OrmarUserDatabase(BaseUserDatabase[UD]): if oauth_accounts and self.oauth_account_model: await model.oauth_accounts.clear(keep_reversed=False) await self._create_oauth_models(model=model, oauth_accounts=oauth_accounts) - user = await self._get_user(id=user.id) - return user + user_db = await self._get_user(id=user.id) + return cast(UD, user_db) async def delete(self, user: UD) -> None: await self.model.objects.delete(id=user.id) async def _create_oauth_models( - self, model: ormar.Model, oauth_accounts: List[BaseOAuthAccount] + self, model: OrmarBaseUserModel, oauth_accounts: List[BaseOAuthAccount] ): - await self.oauth_account_model.objects.bulk_create( - [ + if self.oauth_account_model: + oauth_accounts_db: List[ormar.Model] = [ self.oauth_account_model(user=model, **oacc.dict()) for oacc in oauth_accounts ] - ) + await self.oauth_account_model.objects.bulk_create(oauth_accounts_db) - async def _get_db_user(self, **kwargs: Any) -> Optional[ormar.Model]: + async def _get_db_user(self, **kwargs: Any) -> OrmarBaseUserModel: query = self.model.objects.filter(**kwargs) if self.oauth_account_model is not None: query = query.select_related("oauth_accounts") - return await query.get() + return cast(OrmarBaseUserModel, await query.get()) async def _get_user(self, **kwargs: Any) -> Optional[UD]: try: diff --git a/fastapi_users/db/tortoise.py b/fastapi_users/db/tortoise.py index 2f3021aa..3c17481c 100644 --- a/fastapi_users/db/tortoise.py +++ b/fastapi_users/db/tortoise.py @@ -3,6 +3,7 @@ from typing import Optional, Type from pydantic import UUID4 from tortoise import fields, models from tortoise.exceptions import DoesNotExist +from tortoise.queryset import QuerySetSingle from fastapi_users.db.base import BaseUserDatabase from fastapi_users.models import UD @@ -93,7 +94,7 @@ class TortoiseUserDatabase(BaseUserDatabase[UD]): async def get_by_oauth_account(self, oauth: str, account_id: str) -> Optional[UD]: try: - query = self.model.get( + query: QuerySetSingle[TortoiseBaseUserModel] = self.model.get( oauth_accounts__oauth_name=oauth, oauth_accounts__account_id=account_id ).prefetch_related("oauth_accounts") @@ -132,7 +133,7 @@ class TortoiseUserDatabase(BaseUserDatabase[UD]): await model.save() if oauth_accounts and self.oauth_account_model: - await model.oauth_accounts.all().delete() + await model.oauth_accounts.all().delete() # type: ignore oauth_account_objects = [] for oauth_account in oauth_accounts: oauth_account_objects.append(