Add function get_value to all built-in storage implementations, FSMContext and SceneWizard (#1431)

CHANGES/1431.feature.rst

@@ -0,0 +1 @@
+Add function ``get_value`` to all built-in storage implementations, ``FSMContext`` and ``SceneWizard``

aiogram/fsm/context.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, overload
 
 from aiogram.fsm.storage.base import BaseStorage, StateType, StorageKey
 
@@ -20,6 +20,15 @@ class FSMContext:
     async def get_data(self) -> Dict[str, Any]:
         return await self.storage.get_data(key=self.key)
 
+    @overload
+    async def get_value(self, key: str) -> Any | None: ...
+
+    @overload
+    async def get_value(self, key: str, default: Any) -> Any: ...
+
+    async def get_value(self, key, default=None):
+        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
     ) -> Dict[str, Any]:

aiogram/fsm/scene.py

@@ -4,7 +4,7 @@ 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
+from typing import Any, ClassVar, Dict, List, Optional, Tuple, Type, Union, overload
 
 from typing_extensions import Self
 
@@ -572,6 +572,32 @@ class SceneWizard:
         """
         return await self.state.get_data()
 
+    @overload
+    async def get_value(self, key: str) -> Optional[Any]:
+        """
+        This method returns the value from key in the data of the current state.
+
+        :param key: The keyname of the item you want to return the value from.
+
+        :return: A dictionary containing the data stored in the scene state.
+        """
+        pass
+
+    @overload
+    async def get_value(self, key: str, default: Any) -> Any:
+        """
+        This method returns the value from key in the data of the current state.
+
+        :param key: The keyname of the item you want to return the value from.
+        :param default: Default value to return, if ``key`` was not found.
+
+        :return: A dictionary containing the data stored in the scene state.
+        """
+        pass
+
+    async def get_value(self, key, default=None):
+        return await self.state.get_value(key, default)
+
     async def update_data(
         self, data: Optional[Dict[str, Any]] = None, **kwargs: Any
     ) -> Dict[str, Any]:

aiogram/fsm/storage/base.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
-from typing import Any, AsyncGenerator, Dict, Literal, Optional, Union
+from typing import Any, AsyncGenerator, Dict, Literal, Optional, Union, overload
 
 from aiogram.fsm.state import State
 
@@ -144,6 +144,31 @@ class BaseStorage(ABC):
         """
         pass
 
+    @overload
+    @abstractmethod
+    async def get_value(self, storage_key: StorageKey, dict_key: str) -> Optional[Any]:
+        """
+        Get single value from data by key
+
+        :param storage_key: storage key
+        :param dict_key: value key
+        :return: value stored in key of dict or ``None``
+        """
+        pass
+
+    @overload
+    @abstractmethod
+    async def get_value(self, storage_key: StorageKey, dict_key: str, default: Any) -> Any:
+        """
+        Get single value from data by key
+
+        :param storage_key: storage key
+        :param dict_key: value key
+        :param default: default value to return
+        :return: value stored in key of dict or default
+        """
+        pass
+
     async def update_data(self, key: StorageKey, data: Dict[str, Any]) -> Dict[str, Any]:
         """
         Update date in the storage for key (like dict.update)

aiogram/fsm/storage/memory.py

@@ -1,8 +1,9 @@
 from asyncio import Lock
 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
+from typing import Any, AsyncGenerator, DefaultDict, Dict, Hashable, Optional, overload
 
 from aiogram.fsm.state import State
 from aiogram.fsm.storage.base import (
@@ -49,6 +50,15 @@ class MemoryStorage(BaseStorage):
     async def get_data(self, key: StorageKey) -> Dict[str, Any]:
         return self.storage[key].data.copy()
 
+    @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, dict_key: str, default=None):
+        return copy(self.storage[storage_key].data.get(dict_key, default))
+
 
 class DisabledEventIsolation(BaseEventIsolation):
     @asynccontextmanager

aiogram/fsm/storage/mongo.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional, cast
+from typing import Any, Dict, Optional, cast, overload
 
 from motor.motor_asyncio import AsyncIOMotorClient
 
@@ -115,6 +115,15 @@ 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, dict_key, default=None):
+        return (await self.get_data(storage_key)).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()}

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, Optional, cast, overload
 
 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,
@@ -127,6 +128,19 @@ 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, dict_key, default=None):
+        return (await self.get_data(storage_key)).get(dict_key, default)
+
 
 class RedisEventIsolation(BaseEventIsolation):
     def __init__(

tests/test_fsm/storage/test_storages.py

@@ -22,11 +22,27 @@ class TestStorages:
 
     async def test_set_data(self, storage: BaseStorage, storage_key: StorageKey):
         assert await storage.get_data(key=storage_key) == {}
+        assert await storage.get_value(storage_key=storage_key, dict_key="foo") is None
+        assert (
+            await storage.get_value(storage_key=storage_key, dict_key="foo", default="baz")
+            == "baz"
+        )
 
         await storage.set_data(key=storage_key, data={"foo": "bar"})
         assert await storage.get_data(key=storage_key) == {"foo": "bar"}
+        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"
+        )
+
         await storage.set_data(key=storage_key, data={})
         assert await storage.get_data(key=storage_key) == {}
+        assert await storage.get_value(storage_key=storage_key, dict_key="foo") is None
+        assert (
+            await storage.get_value(storage_key=storage_key, dict_key="foo", default="baz")
+            == "baz"
+        )
 
     async def test_update_data(self, storage: BaseStorage, storage_key: StorageKey):
         assert await storage.get_data(key=storage_key) == {}

tests/test_fsm/test_context.py

@@ -34,6 +34,10 @@ class TestFSMContext:
         assert await state2.get_data() == {}
         assert await state3.get_data() == {}
 
+        assert await state.get_value("foo") == "bar"
+        assert await state2.get_value("foo") is None
+        assert await state3.get_value("foo", "baz") == "baz"
+
         await state2.set_state("experiments")
         assert await state.get_state() == "test"
         assert await state3.get_state() is None

tests/test_fsm/test_scene.py

@@ -1004,6 +1004,24 @@ class TestSceneWizard:
 
         wizard.state.get_data.assert_called_once_with()
 
+    async def test_scene_wizard_get_value_with_default(self):
+        wizard = SceneWizard(
+            scene_config=AsyncMock(),
+            manager=AsyncMock(),
+            state=AsyncMock(),
+            update_type="message",
+            event=AsyncMock(),
+            data={},
+        )
+        args = ("test_key", "test_default")
+        value = "test_value"
+        wizard.state.get_value = AsyncMock(return_value=value)
+
+        result = await wizard.get_value(*args)
+        wizard.state.get_value.assert_called_once_with(*args)
+
+        assert result == value
+
     async def test_scene_wizard_update_data_if_data(self):
         wizard = SceneWizard(
             scene_config=AsyncMock(),