Initial commit

venv/Lib/site-packages/aiogram/fsm/storage/base.py

@@ -0,0 +1,212 @@
+from abc import ABC, abstractmethod
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import Any, AsyncGenerator, Dict, Literal, Optional, Union, overload
+
+from aiogram.fsm.state import State
+
+StateType = Optional[Union[str, State]]
+
+DEFAULT_DESTINY = "default"
+
+
+@dataclass(frozen=True)
+class StorageKey:
+    bot_id: int
+    chat_id: int
+    user_id: int
+    thread_id: Optional[int] = None
+    business_connection_id: Optional[str] = None
+    destiny: str = DEFAULT_DESTINY
+
+
+class KeyBuilder(ABC):
+    """Base class for key builder."""
+
+    @abstractmethod
+    def build(
+        self,
+        key: StorageKey,
+        part: Optional[Literal["data", "state", "lock"]] = None,
+    ) -> str:
+        """
+        Build key to be used in storage's db queries
+
+        :param key: contextual key
+        :param part: part of the record
+        :return: key to be used in storage's db queries
+        """
+        pass
+
+
+class DefaultKeyBuilder(KeyBuilder):
+    """
+    Simple key builder with default prefix.
+
+    Generates a colon-joined string with prefix, chat_id, user_id,
+    optional bot_id, business_connection_id, destiny and field.
+
+    Format:
+     :code:`<prefix>:<bot_id?>:<business_connection_id?>:<chat_id>:<user_id>:<destiny?>:<field?>`
+    """
+
+    def __init__(
+        self,
+        *,
+        prefix: str = "fsm",
+        separator: str = ":",
+        with_bot_id: bool = False,
+        with_business_connection_id: bool = False,
+        with_destiny: bool = False,
+    ) -> None:
+        """
+        :param prefix: prefix for all records
+        :param separator: separator
+        :param with_bot_id: include Bot id in the key
+        :param with_business_connection_id: include business connection id
+        :param with_destiny: include destiny key
+        """
+        self.prefix = prefix
+        self.separator = separator
+        self.with_bot_id = with_bot_id
+        self.with_business_connection_id = with_business_connection_id
+        self.with_destiny = with_destiny
+
+    def build(
+        self,
+        key: StorageKey,
+        part: Optional[Literal["data", "state", "lock"]] = None,
+    ) -> str:
+        parts = [self.prefix]
+        if self.with_bot_id:
+            parts.append(str(key.bot_id))
+        if self.with_business_connection_id and key.business_connection_id:
+            parts.append(str(key.business_connection_id))
+        parts.append(str(key.chat_id))
+        if key.thread_id:
+            parts.append(str(key.thread_id))
+        parts.append(str(key.user_id))
+        if self.with_destiny:
+            parts.append(key.destiny)
+        elif key.destiny != DEFAULT_DESTINY:
+            error_message = (
+                "Default key builder is not configured to use key destiny other than the default."
+                "\n\nProbably, you should set `with_destiny=True` in for DefaultKeyBuilder."
+            )
+            raise ValueError(error_message)
+        if part:
+            parts.append(part)
+        return self.separator.join(parts)
+
+
+class BaseStorage(ABC):
+    """
+    Base class for all FSM storages
+    """
+
+    @abstractmethod
+    async def set_state(self, key: StorageKey, state: StateType = None) -> None:
+        """
+        Set state for specified key
+
+        :param key: storage key
+        :param state: new state
+        """
+        pass
+
+    @abstractmethod
+    async def get_state(self, key: StorageKey) -> Optional[str]:
+        """
+        Get key state
+
+        :param key: storage key
+        :return: current state
+        """
+        pass
+
+    @abstractmethod
+    async def set_data(self, key: StorageKey, data: Dict[str, Any]) -> None:
+        """
+        Write data (replace)
+
+        :param key: storage key
+        :param data: new data
+        """
+        pass
+
+    @abstractmethod
+    async def get_data(self, key: StorageKey) -> Dict[str, Any]:
+        """
+        Get current data for key
+
+        :param key: storage key
+        :return: current data
+        """
+        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)
+
+        :param key: storage key
+        :param data: partial data
+        :return: new data
+        """
+        current_data = await self.get_data(key=key)
+        current_data.update(data)
+        await self.set_data(key=key, data=current_data)
+        return current_data.copy()
+
+    @abstractmethod
+    async def close(self) -> None:  # pragma: no cover
+        """
+        Close storage (database connection, file or etc.)
+        """
+        pass
+
+
+class BaseEventIsolation(ABC):
+    @abstractmethod
+    @asynccontextmanager
+    async def lock(self, key: StorageKey) -> AsyncGenerator[None, None]:
+        """
+        Isolate events with lock.
+        Will be used as context manager
+
+        :param key: storage key
+        :return: An async generator
+        """
+        yield None
+
+    @abstractmethod
+    async def close(self) -> None:
+        pass

venv/Lib/site-packages/aiogram/fsm/storage/memory.py

@@ -0,0 +1,87 @@
+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, overload
+
+from aiogram.fsm.state import State
+from aiogram.fsm.storage.base import (
+    BaseEventIsolation,
+    BaseStorage,
+    StateType,
+    StorageKey,
+)
+
+
+@dataclass
+class MemoryStorageRecord:
+    data: Dict[str, Any] = field(default_factory=dict)
+    state: Optional[str] = None
+
+
+class MemoryStorage(BaseStorage):
+    """
+    Default FSM storage, stores all data in :class:`dict` and loss everything on shutdown
+
+    .. warning::
+
+        Is not recommended using in production in due to you will lose all data
+        when your bot restarts
+    """
+
+    def __init__(self) -> None:
+        self.storage: DefaultDict[StorageKey, MemoryStorageRecord] = defaultdict(
+            MemoryStorageRecord
+        )
+
+    async def close(self) -> None:
+        pass
+
+    async def set_state(self, key: StorageKey, state: StateType = None) -> None:
+        self.storage[key].state = state.state if isinstance(state, State) else state
+
+    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:
+        self.storage[key].data = data.copy()
+
+    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
+    async def lock(self, key: StorageKey) -> AsyncGenerator[None, None]:
+        yield
+
+    async def close(self) -> None:
+        pass
+
+
+class SimpleEventIsolation(BaseEventIsolation):
+    def __init__(self) -> None:
+        # TODO: Unused locks cleaner is needed
+        self._locks: DefaultDict[Hashable, Lock] = defaultdict(Lock)
+
+    @asynccontextmanager
+    async def lock(self, key: StorageKey) -> AsyncGenerator[None, None]:
+        lock = self._locks[key]
+        async with lock:
+            yield
+
+    async def close(self) -> None:
+        self._locks.clear()

venv/Lib/site-packages/aiogram/fsm/storage/mongo.py

@@ -0,0 +1,130 @@
+from typing import Any, Dict, Optional, cast
+
+from motor.motor_asyncio import AsyncIOMotorClient
+
+from aiogram.fsm.state import State
+from aiogram.fsm.storage.base import (
+    BaseStorage,
+    DefaultKeyBuilder,
+    KeyBuilder,
+    StateType,
+    StorageKey,
+)
+
+
+class MongoStorage(BaseStorage):
+    """
+    MongoDB storage required :code:`motor` package installed (:code:`pip install motor`)
+    """
+
+    def __init__(
+        self,
+        client: AsyncIOMotorClient,
+        key_builder: Optional[KeyBuilder] = None,
+        db_name: str = "aiogram_fsm",
+        collection_name: str = "states_and_data",
+    ) -> None:
+        """
+        :param client: Instance of AsyncIOMotorClient
+        :param key_builder: builder that helps to convert contextual key to string
+        :param db_name: name of the MongoDB database for FSM
+        :param collection_name: name of the collection for storing FSM states and data
+        """
+        if key_builder is None:
+            key_builder = DefaultKeyBuilder()
+        self._client = client
+        self._database = self._client[db_name]
+        self._collection = self._database[collection_name]
+        self._key_builder = key_builder
+
+    @classmethod
+    def from_url(
+        cls, url: str, connection_kwargs: Optional[Dict[str, Any]] = None, **kwargs: Any
+    ) -> "MongoStorage":
+        """
+        Create an instance of :class:`MongoStorage` with specifying the connection string
+
+        :param url: for example :code:`mongodb://user:password@host:port`
+        :param connection_kwargs: see :code:`motor` docs
+        :param kwargs: arguments to be passed to :class:`MongoStorage`
+        :return: an instance of :class:`MongoStorage`
+        """
+        if connection_kwargs is None:
+            connection_kwargs = {}
+        client = AsyncIOMotorClient(url, **connection_kwargs)
+        return cls(client=client, **kwargs)
+
+    async def close(self) -> None:
+        """Cleanup client resources and disconnect from MongoDB."""
+        self._client.close()
+
+    def resolve_state(self, value: StateType) -> Optional[str]:
+        if value is None:
+            return None
+        if isinstance(value, State):
+            return value.state
+        return str(value)
+
+    async def set_state(self, key: StorageKey, state: StateType = None) -> None:
+        document_id = self._key_builder.build(key)
+        if state is None:
+            updated = await self._collection.find_one_and_update(
+                filter={"_id": document_id},
+                update={"$unset": {"state": 1}},
+                projection={"_id": 0},
+                return_document=True,
+            )
+            if updated == {}:
+                await self._collection.delete_one({"_id": document_id})
+        else:
+            await self._collection.update_one(
+                filter={"_id": document_id},
+                update={"$set": {"state": self.resolve_state(state)}},
+                upsert=True,
+            )
+
+    async def get_state(self, key: StorageKey) -> Optional[str]:
+        document_id = self._key_builder.build(key)
+        document = await self._collection.find_one({"_id": document_id})
+        if document is None:
+            return None
+        return document.get("state")
+
+    async def set_data(self, key: StorageKey, data: Dict[str, Any]) -> None:
+        document_id = self._key_builder.build(key)
+        if not data:
+            updated = await self._collection.find_one_and_update(
+                filter={"_id": document_id},
+                update={"$unset": {"data": 1}},
+                projection={"_id": 0},
+                return_document=True,
+            )
+            if updated == {}:
+                await self._collection.delete_one({"_id": document_id})
+        else:
+            await self._collection.update_one(
+                filter={"_id": document_id},
+                update={"$set": {"data": data}},
+                upsert=True,
+            )
+
+    async def get_data(self, key: StorageKey) -> Dict[str, Any]:
+        document_id = self._key_builder.build(key)
+        document = await self._collection.find_one({"_id": document_id})
+        if document is None or not document.get("data"):
+            return {}
+        return cast(Dict[str, Any], document["data"])
+
+    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()}
+        update_result = await self._collection.find_one_and_update(
+            filter={"_id": document_id},
+            update={"$set": update_with},
+            upsert=True,
+            return_document=True,
+            projection={"_id": 0},
+        )
+        if not update_result:
+            await self._collection.delete_one({"_id": document_id})
+        return update_result.get("data", {})

venv/Lib/site-packages/aiogram/fsm/storage/redis.py

@@ -0,0 +1,169 @@
+import json
+from contextlib import asynccontextmanager
+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.state import State
+from aiogram.fsm.storage.base import (
+    BaseEventIsolation,
+    BaseStorage,
+    DefaultKeyBuilder,
+    KeyBuilder,
+    StateType,
+    StorageKey,
+)
+
+DEFAULT_REDIS_LOCK_KWARGS = {"timeout": 60}
+_JsonLoads = Callable[..., Any]
+_JsonDumps = Callable[..., str]
+
+
+class RedisStorage(BaseStorage):
+    """
+    Redis storage required :code:`redis` package installed (:code:`pip install redis`)
+    """
+
+    def __init__(
+        self,
+        redis: Redis,
+        key_builder: Optional[KeyBuilder] = None,
+        state_ttl: Optional[ExpiryT] = None,
+        data_ttl: Optional[ExpiryT] = None,
+        json_loads: _JsonLoads = json.loads,
+        json_dumps: _JsonDumps = json.dumps,
+    ) -> None:
+        """
+        :param redis: Instance of Redis connection
+        :param key_builder: builder that helps to convert contextual key to string
+        :param state_ttl: TTL for state records
+        :param data_ttl: TTL for data records
+        """
+        if key_builder is None:
+            key_builder = DefaultKeyBuilder()
+        self.redis = redis
+        self.key_builder = key_builder
+        self.state_ttl = state_ttl
+        self.data_ttl = data_ttl
+        self.json_loads = json_loads
+        self.json_dumps = json_dumps
+
+    @classmethod
+    def from_url(
+        cls, url: str, connection_kwargs: Optional[Dict[str, Any]] = None, **kwargs: Any
+    ) -> "RedisStorage":
+        """
+        Create an instance of :class:`RedisStorage` with specifying the connection string
+
+        :param url: for example :code:`redis://user:password@host:port/db`
+        :param connection_kwargs: see :code:`redis` docs
+        :param kwargs: arguments to be passed to :class:`RedisStorage`
+        :return: an instance of :class:`RedisStorage`
+        """
+        if connection_kwargs is None:
+            connection_kwargs = {}
+        pool = ConnectionPool.from_url(url, **connection_kwargs)
+        redis = Redis(connection_pool=pool)
+        return cls(redis=redis, **kwargs)
+
+    def create_isolation(self, **kwargs: Any) -> "RedisEventIsolation":
+        return RedisEventIsolation(redis=self.redis, key_builder=self.key_builder, **kwargs)
+
+    async def close(self) -> None:
+        await self.redis.aclose(close_connection_pool=True)
+
+    async def set_state(
+        self,
+        key: StorageKey,
+        state: StateType = None,
+    ) -> None:
+        redis_key = self.key_builder.build(key, "state")
+        if state is None:
+            await self.redis.delete(redis_key)
+        else:
+            await self.redis.set(
+                redis_key,
+                cast(str, state.state if isinstance(state, State) else state),
+                ex=self.state_ttl,
+            )
+
+    async def get_state(
+        self,
+        key: StorageKey,
+    ) -> Optional[str]:
+        redis_key = self.key_builder.build(key, "state")
+        value = await self.redis.get(redis_key)
+        if isinstance(value, bytes):
+            return value.decode("utf-8")
+        return cast(Optional[str], value)
+
+    async def set_data(
+        self,
+        key: StorageKey,
+        data: Dict[str, Any],
+    ) -> None:
+        redis_key = self.key_builder.build(key, "data")
+        if not data:
+            await self.redis.delete(redis_key)
+            return
+        await self.redis.set(
+            redis_key,
+            self.json_dumps(data),
+            ex=self.data_ttl,
+        )
+
+    async def get_data(
+        self,
+        key: StorageKey,
+    ) -> Dict[str, Any]:
+        redis_key = self.key_builder.build(key, "data")
+        value = await self.redis.get(redis_key)
+        if value is None:
+            return {}
+        if isinstance(value, bytes):
+            value = value.decode("utf-8")
+        return cast(Dict[str, Any], self.json_loads(value))
+
+
+class RedisEventIsolation(BaseEventIsolation):
+    def __init__(
+        self,
+        redis: Redis,
+        key_builder: Optional[KeyBuilder] = None,
+        lock_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        if key_builder is None:
+            key_builder = DefaultKeyBuilder()
+        if lock_kwargs is None:
+            lock_kwargs = DEFAULT_REDIS_LOCK_KWARGS
+        self.redis = redis
+        self.key_builder = key_builder
+        self.lock_kwargs = lock_kwargs
+
+    @classmethod
+    def from_url(
+        cls,
+        url: str,
+        connection_kwargs: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> "RedisEventIsolation":
+        if connection_kwargs is None:
+            connection_kwargs = {}
+        pool = ConnectionPool.from_url(url, **connection_kwargs)
+        redis = Redis(connection_pool=pool)
+        return cls(redis=redis, **kwargs)
+
+    @asynccontextmanager
+    async def lock(
+        self,
+        key: StorageKey,
+    ) -> AsyncGenerator[None, None]:
+        redis_key = self.key_builder.build(key, "lock")
+        async with self.redis.lock(name=redis_key, **self.lock_kwargs, lock_class=Lock):
+            yield None
+
+    async def close(self) -> None:
+        pass