From b435e446a55ec427cce82d1afc2eab5c4c0be68d Mon Sep 17 00:00:00 2001 From: Alex Root Junior Date: Fri, 4 Aug 2017 15:59:09 +0300 Subject: [PATCH] Split FSM storages to different modules. --- aiogram/contrib/fsm_storage/__init__.py | 0 aiogram/contrib/fsm_storage/memory.py | 78 +++++++++++++++++++++++++ aiogram/dispatcher/storage.py | 75 ------------------------ 3 files changed, 78 insertions(+), 75 deletions(-) create mode 100644 aiogram/contrib/fsm_storage/__init__.py create mode 100644 aiogram/contrib/fsm_storage/memory.py diff --git a/aiogram/contrib/fsm_storage/__init__.py b/aiogram/contrib/fsm_storage/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/aiogram/contrib/fsm_storage/memory.py b/aiogram/contrib/fsm_storage/memory.py new file mode 100644 index 00000000..cbd8f243 --- /dev/null +++ b/aiogram/contrib/fsm_storage/memory.py @@ -0,0 +1,78 @@ +import typing + +from aiogram.dispatcher import BaseStorage + + +class MemoryStorage(BaseStorage): + """ + In-memory based states storage. + + This type of storage is not recommended for usage in bots, because you will lost all states after restarting. + """ + + def __init__(self): + self.data = {} + + def _get_chat(self, chat_id): + chat_id = str(chat_id) + if chat_id not in self.data: + self.data[chat_id] = {} + return self.data[chat_id] + + def _get_user(self, chat_id, user_id): + chat = self._get_chat(chat_id) + chat_id = str(chat_id) + user_id = str(user_id) + if user_id not in self.data[chat_id]: + self.data[chat_id][user_id] = {'state': None, 'data': {}} + return self.data[chat_id][user_id] + + async def get_state(self, *, + chat: typing.Union[str, int, None] = None, + user: typing.Union[str, int, None] = None, + default: typing.Optional[str] = None) -> typing.Optional[str]: + chat, user = self.check_address(chat=chat, user=user) + user = self._get_user(chat, user) + return user['state'] + + async def get_data(self, *, + chat: typing.Union[str, int, None] = None, + user: typing.Union[str, int, None] = None, + default: typing.Optional[str] = None) -> typing.Dict: + chat, user = self.check_address(chat=chat, user=user) + user = self._get_user(chat, user) + return user['data'] + + async def update_data(self, *, + chat: typing.Union[str, int, None] = None, + user: typing.Union[str, int, None] = None, + data: typing.Dict = None, **kwargs): + chat, user = self.check_address(chat=chat, user=user) + user = self._get_user(chat, user) + if data is None: + data = [] + user['data'].update(data, **kwargs) + + async def set_state(self, *, + chat: typing.Union[str, int, None] = None, + user: typing.Union[str, int, None] = None, + state: typing.AnyStr = None): + chat, user = self.check_address(chat=chat, user=user) + user = self._get_user(chat, user) + user['state'] = state + + async def set_data(self, *, + chat: typing.Union[str, int, None] = None, + user: typing.Union[str, int, None] = None, + data: typing.Dict = None): + chat, user = self.check_address(chat=chat, user=user) + user = self._get_user(chat, user) + user['data'] = data + + async def reset_state(self, *, + chat: typing.Union[str, int, None] = None, + user: typing.Union[str, int, None] = None, + with_data: typing.Optional[bool] = True): + await self.set_state(chat=chat, user=user, state=None) + if with_data: + await self.set_data(chat=chat, user=user, data={}) diff --git a/aiogram/dispatcher/storage.py b/aiogram/dispatcher/storage.py index e33a0ae5..4deb3b54 100644 --- a/aiogram/dispatcher/storage.py +++ b/aiogram/dispatcher/storage.py @@ -238,78 +238,3 @@ class DisabledStorage(BaseStorage): user: typing.Union[str, int, None] = None, data: typing.Dict = None): pass - - -class MemoryStorage(BaseStorage): - """ - In-memory based states storage. - - This type of storage is not recommended for usage in bots, because you will lost all states after restarting. - """ - - def __init__(self): - self.data = {} - - def _get_chat(self, chat_id): - chat_id = str(chat_id) - if chat_id not in self.data: - self.data[chat_id] = {} - return self.data[chat_id] - - def _get_user(self, chat_id, user_id): - chat = self._get_chat(chat_id) - chat_id = str(chat_id) - user_id = str(user_id) - if user_id not in self.data[chat_id]: - self.data[chat_id][user_id] = {'state': None, 'data': {}} - return self.data[chat_id][user_id] - - async def get_state(self, *, - chat: typing.Union[str, int, None] = None, - user: typing.Union[str, int, None] = None, - default: typing.Optional[str] = None) -> typing.Optional[str]: - chat, user = self.check_address(chat=chat, user=user) - user = self._get_user(chat, user) - return user['state'] - - async def get_data(self, *, - chat: typing.Union[str, int, None] = None, - user: typing.Union[str, int, None] = None, - default: typing.Optional[str] = None) -> typing.Dict: - chat, user = self.check_address(chat=chat, user=user) - user = self._get_user(chat, user) - return user['data'] - - async def update_data(self, *, - chat: typing.Union[str, int, None] = None, - user: typing.Union[str, int, None] = None, - data: typing.Dict = None, **kwargs): - chat, user = self.check_address(chat=chat, user=user) - user = self._get_user(chat, user) - if data is None: - data = [] - user['data'].update(data, **kwargs) - - async def set_state(self, *, - chat: typing.Union[str, int, None] = None, - user: typing.Union[str, int, None] = None, - state: typing.AnyStr = None): - chat, user = self.check_address(chat=chat, user=user) - user = self._get_user(chat, user) - user['state'] = state - - async def set_data(self, *, - chat: typing.Union[str, int, None] = None, - user: typing.Union[str, int, None] = None, - data: typing.Dict = None): - chat, user = self.check_address(chat=chat, user=user) - user = self._get_user(chat, user) - user['data'] = data - - async def reset_state(self, *, - chat: typing.Union[str, int, None] = None, - user: typing.Union[str, int, None] = None, - with_data: typing.Optional[bool] = True): - await self.set_state(chat=chat, user=user, state=None) - if with_data: - await self.set_data(chat=chat, user=user, data={})