Implement get_value method in BaseStorage and remove redundant implementations

This commit is contained in:
Arthur Khachaturov 2024-10-23 00:55:27 +03:00
parent e5f1d5be60
commit 2999e8d62d
No known key found for this signature in database
GPG key ID: CAC2B7EB6DF45D55
3 changed files with 5 additions and 33 deletions

View file

@ -167,10 +167,11 @@ class BaseStorage(ABC):
"""
pass
@abstractmethod
async def get_value(
self, storage_key: StorageKey, dict_key: str, default: Optional[Any] = None
) -> Optional[Any]: ...
) -> Optional[Any]:
data = await self.get_data(storage_key)
return data.get(dict_key, default)
async def update_data(self, key: StorageKey, data: Dict[str, Any]) -> Dict[str, Any]:
"""

View file

@ -1,4 +1,4 @@
from typing import Any, Dict, Optional, cast, overload
from typing import Any, Dict, Optional, cast
from motor.motor_asyncio import AsyncIOMotorClient
@ -115,18 +115,6 @@ class MongoStorage(BaseStorage):
return {}
return cast(Dict[str, Any], document["data"])
@overload
async def get_value(self, storage_key: StorageKey, dict_key: str) -> Optional[Any]: ...
@overload
async def get_value(self, storage_key: StorageKey, dict_key: str, default: Any) -> Any: ...
async def get_value(
self, storage_key: StorageKey, dict_key: str, default: Optional[Any] = None
) -> Optional[Any]:
data = await self.get_data(storage_key)
return data.get(dict_key, default)
async def update_data(self, key: StorageKey, data: Dict[str, Any]) -> Dict[str, Any]:
document_id = self._key_builder.build(key)
update_with = {f"data.{key}": value for key, value in data.items()}

View file

@ -1,13 +1,12 @@
import json
from contextlib import asynccontextmanager
from typing import Any, AsyncGenerator, Callable, Dict, Optional, cast, overload
from typing import Any, AsyncGenerator, Callable, Dict, Optional, cast
from redis.asyncio.client import Redis
from redis.asyncio.connection import ConnectionPool
from redis.asyncio.lock import Lock
from redis.typing import ExpiryT
from aiogram.fsm import storage
from aiogram.fsm.state import State
from aiogram.fsm.storage.base import (
BaseEventIsolation,
@ -128,22 +127,6 @@ class RedisStorage(BaseStorage):
value = value.decode("utf-8")
return cast(Dict[str, Any], self.json_loads(value))
@overload
async def get_value(
self,
storage_key: StorageKey,
dict_key: str,
) -> Optional[Any]: ...
@overload
async def get_value(self, storage_key: StorageKey, dict_key: str, default: Any) -> Any: ...
async def get_value(
self, storage_key: StorageKey, dict_key: str, default: Optional[Any] = None
) -> Optional[Any]:
data = await self.get_data(storage_key)
return data.get(dict_key, default)
class RedisEventIsolation(BaseEventIsolation):
def __init__(