Refactor FSM Storage methods input types to calm down MyPy. (#1683)

* Refactor methods input types to calm down MyPy.
- `FSMContext.set_data`
- `FSMContext.update_data`
- `BaseStorage.set_data`
- `BaseStorage.update_data`
- `BaseStorage`'s child methods
- `SceneWizard.set_data`
- `SceneWizard.update_data`

* Add 1683.feature.rst

* Remove re-init in `DataNotDictLikeError`

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`

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.
+    """

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)

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
         """

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)
 

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]:

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(

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)

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"}