diff --git a/CHANGES/1683.feature.rst b/CHANGES/1683.feature.rst new file mode 100644 index 00000000..97a53e2e --- /dev/null +++ b/CHANGES/1683.feature.rst @@ -0,0 +1,11 @@ +Refactor methods input types to calm down MyPy. #1682 + +`Dict[str, Any]` is replaced with `Mapping[str, Any]` in the following methods: + +- `FSMContext.set_data` +- `FSMContext.update_data` +- `BaseStorage.set_data` +- `BaseStorage.update_data` +- `BaseStorage's child methods` +- `SceneWizard.set_data` +- `SceneWizard.update_data` diff --git a/aiogram/exceptions.py b/aiogram/exceptions.py index d195aa7b..a1b73d56 100644 --- a/aiogram/exceptions.py +++ b/aiogram/exceptions.py @@ -197,3 +197,9 @@ class ClientDecodeError(AiogramError): f"{original_type.__module__}.{original_type.__name__}: {self.original}\n" f"Content: {self.data}" ) + + +class DataNotDictLikeError(DetailedAiogramError): + """ + Exception raised when data is not dict-like. + """ diff --git a/aiogram/fsm/context.py b/aiogram/fsm/context.py index c9432c11..f353df5d 100644 --- a/aiogram/fsm/context.py +++ b/aiogram/fsm/context.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, Optional, overload +from typing import Any, Dict, Mapping, Optional, overload from aiogram.fsm.storage.base import BaseStorage, StateType, StorageKey @@ -14,7 +14,7 @@ class FSMContext: async def get_state(self) -> Optional[str]: return await self.storage.get_state(key=self.key) - async def set_data(self, data: Dict[str, Any]) -> None: + async def set_data(self, data: Mapping[str, Any]) -> None: await self.storage.set_data(key=self.key, data=data) async def get_data(self) -> Dict[str, Any]: @@ -30,7 +30,7 @@ class FSMContext: return await self.storage.get_value(storage_key=self.key, dict_key=key, default=default) async def update_data( - self, data: Optional[Dict[str, Any]] = None, **kwargs: Any + self, data: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: if data: kwargs.update(data) diff --git a/aiogram/fsm/scene.py b/aiogram/fsm/scene.py index a84acf04..94def0d6 100644 --- a/aiogram/fsm/scene.py +++ b/aiogram/fsm/scene.py @@ -4,7 +4,18 @@ import inspect from collections import defaultdict from dataclasses import dataclass, replace from enum import Enum, auto -from typing import Any, ClassVar, Dict, List, Optional, Tuple, Type, Union, overload +from typing import ( + Any, + ClassVar, + Dict, + List, + Mapping, + Optional, + Tuple, + Type, + Union, + overload, +) from typing_extensions import Self @@ -577,11 +588,11 @@ class SceneWizard: await action_config[event_type].call(self.scene, self.event, **{**self.data, **kwargs}) return True - async def set_data(self, data: Dict[str, Any]) -> None: + async def set_data(self, data: Mapping[str, Any]) -> None: """ Sets custom data in the current state. - :param data: A dictionary containing the custom data to be set in the current state. + :param data: A mapping containing the custom data to be set in the current state. :return: None """ await self.state.set_data(data=data) @@ -621,12 +632,12 @@ class SceneWizard: return await self.state.get_value(key, default) async def update_data( - self, data: Optional[Dict[str, Any]] = None, **kwargs: Any + self, data: Optional[Mapping[str, Any]] = None, **kwargs: Any ) -> Dict[str, Any]: """ This method updates the data stored in the current state - :param data: Optional dictionary of data to update. + :param data: Optional mapping of data to update. :param kwargs: Additional key-value pairs of data to update. :return: Dictionary of updated data """ diff --git a/aiogram/fsm/storage/base.py b/aiogram/fsm/storage/base.py index 96cb2cbe..7a1059d6 100644 --- a/aiogram/fsm/storage/base.py +++ b/aiogram/fsm/storage/base.py @@ -1,7 +1,16 @@ from abc import ABC, abstractmethod from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import Any, AsyncGenerator, Dict, Literal, Optional, Union, overload +from typing import ( + Any, + AsyncGenerator, + Dict, + Literal, + Mapping, + Optional, + Union, + overload, +) from aiogram.fsm.state import State @@ -125,7 +134,7 @@ class BaseStorage(ABC): pass @abstractmethod - async def set_data(self, key: StorageKey, data: Dict[str, Any]) -> None: + async def set_data(self, key: StorageKey, data: Mapping[str, Any]) -> None: """ Write data (replace) @@ -173,7 +182,7 @@ class BaseStorage(ABC): 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]: + async def update_data(self, key: StorageKey, data: Mapping[str, Any]) -> Dict[str, Any]: """ Update date in the storage for key (like dict.update) diff --git a/aiogram/fsm/storage/memory.py b/aiogram/fsm/storage/memory.py index f26d15c0..b5eebe4a 100644 --- a/aiogram/fsm/storage/memory.py +++ b/aiogram/fsm/storage/memory.py @@ -3,8 +3,18 @@ from collections import defaultdict from contextlib import asynccontextmanager from copy import copy from dataclasses import dataclass, field -from typing import Any, AsyncGenerator, DefaultDict, Dict, Hashable, Optional, overload +from typing import ( + Any, + AsyncGenerator, + DefaultDict, + Dict, + Hashable, + Mapping, + Optional, + overload, +) +from aiogram.exceptions import DataNotDictLikeError from aiogram.fsm.state import State from aiogram.fsm.storage.base import ( BaseEventIsolation, @@ -44,7 +54,11 @@ class MemoryStorage(BaseStorage): async def get_state(self, key: StorageKey) -> Optional[str]: return self.storage[key].state - async def set_data(self, key: StorageKey, data: Dict[str, Any]) -> None: + async def set_data(self, key: StorageKey, data: Mapping[str, Any]) -> None: + if not isinstance(data, dict): + raise DataNotDictLikeError( + f"Data must be a dict or dict-like object, got {type(data).__name__}" + ) self.storage[key].data = data.copy() async def get_data(self, key: StorageKey) -> Dict[str, Any]: diff --git a/aiogram/fsm/storage/mongo.py b/aiogram/fsm/storage/mongo.py index b4b1eeaa..b44aeeb9 100644 --- a/aiogram/fsm/storage/mongo.py +++ b/aiogram/fsm/storage/mongo.py @@ -1,7 +1,8 @@ -from typing import Any, Dict, Optional, cast +from typing import Any, Dict, Mapping, Optional, cast from motor.motor_asyncio import AsyncIOMotorClient +from aiogram.exceptions import DataNotDictLikeError from aiogram.fsm.state import State from aiogram.fsm.storage.base import ( BaseStorage, @@ -90,7 +91,12 @@ class MongoStorage(BaseStorage): return None return document.get("state") - async def set_data(self, key: StorageKey, data: Dict[str, Any]) -> None: + async def set_data(self, key: StorageKey, data: Mapping[str, Any]) -> None: + if not isinstance(data, dict): + raise DataNotDictLikeError( + f"Data must be a dict or dict-like object, got {type(data).__name__}" + ) + document_id = self._key_builder.build(key) if not data: updated = await self._collection.find_one_and_update( @@ -115,7 +121,7 @@ class MongoStorage(BaseStorage): return {} return cast(Dict[str, Any], document["data"]) - async def update_data(self, key: StorageKey, data: Dict[str, Any]) -> Dict[str, Any]: + async def update_data(self, key: StorageKey, data: Mapping[str, Any]) -> Dict[str, Any]: document_id = self._key_builder.build(key) update_with = {f"data.{key}": value for key, value in data.items()} update_result = await self._collection.find_one_and_update( diff --git a/aiogram/fsm/storage/redis.py b/aiogram/fsm/storage/redis.py index edcf9d63..4136972d 100644 --- a/aiogram/fsm/storage/redis.py +++ b/aiogram/fsm/storage/redis.py @@ -1,12 +1,13 @@ import json from contextlib import asynccontextmanager -from typing import Any, AsyncGenerator, Callable, Dict, Optional, cast +from typing import Any, AsyncGenerator, Callable, Dict, Mapping, 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.exceptions import DataNotDictLikeError from aiogram.fsm.state import State from aiogram.fsm.storage.base import ( BaseEventIsolation, @@ -103,8 +104,13 @@ class RedisStorage(BaseStorage): async def set_data( self, key: StorageKey, - data: Dict[str, Any], + data: Mapping[str, Any], ) -> None: + if not isinstance(data, dict): + raise DataNotDictLikeError( + f"Data must be a dict or dict-like object, got {type(data).__name__}" + ) + redis_key = self.key_builder.build(key, "data") if not data: await self.redis.delete(redis_key) diff --git a/tests/test_fsm/storage/test_storages.py b/tests/test_fsm/storage/test_storages.py index 64d4d734..44424ed2 100644 --- a/tests/test_fsm/storage/test_storages.py +++ b/tests/test_fsm/storage/test_storages.py @@ -1,5 +1,8 @@ +from typing import TypedDict + import pytest +from aiogram.exceptions import DataNotDictLikeError from aiogram.fsm.storage.base import BaseStorage, StorageKey @@ -44,6 +47,21 @@ class TestStorages: == "baz" ) + class CustomTypedDict(TypedDict, total=False): + foo: str + bar: str + + await storage.set_data(key=storage_key, data=CustomTypedDict(foo="bar", bar="baz")) + assert await storage.get_data(key=storage_key) == {"foo": "bar", "bar": "baz"} + assert await storage.get_value(storage_key=storage_key, dict_key="foo") == "bar" + assert ( + await storage.get_value(storage_key=storage_key, dict_key="foo", default="baz") + == "bar" + ) + + with pytest.raises(DataNotDictLikeError): + await storage.set_data(key=storage_key, data=()) + async def test_update_data(self, storage: BaseStorage, storage_key: StorageKey): assert await storage.get_data(key=storage_key) == {} assert await storage.update_data(key=storage_key, data={"foo": "bar"}) == {"foo": "bar"}