Add function get_value to BaseStorage (#1594)

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

* Fix type hints

* Split up return statements in `get_value` functions

* Implement `get_value` method in `BaseStorage` and remove redundant implementations

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) -> Optional[Any]: ...
+
+    @overload
+    async def get_value(self, key: str, default: Any) -> Any: ...
+
+    async def get_value(self, key: str, default: Optional[Any] = None) -> Optional[Any]:
+        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: str, default: Optional[Any] = None) -> Optional[Any]:
+        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,35 @@ class BaseStorage(ABC):
         """
         pass
 
+    @overload
+    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
+    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 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]:
         """
         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,18 @@ 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: StorageKey, dict_key: str, default: Optional[Any] = None
+    ) -> Optional[Any]:
+        data = self.storage[storage_key].data
+        return copy(data.get(dict_key, default))
+
 
 class DisabledEventIsolation(BaseEventIsolation):
     @asynccontextmanager

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(),