From 1413baf4b84189b3abc482de5033b5f2cdc1f76f Mon Sep 17 00:00:00 2001 From: Alex Root Junior Date: Tue, 20 Jun 2017 17:54:26 +0300 Subject: [PATCH] Provide async storage. --- aiogram/dispatcher/state.py | 303 +++++++++++++++++++++++++++++++++++- 1 file changed, 300 insertions(+), 3 deletions(-) diff --git a/aiogram/dispatcher/state.py b/aiogram/dispatcher/state.py index a2f087a5..ea543fd8 100644 --- a/aiogram/dispatcher/state.py +++ b/aiogram/dispatcher/state.py @@ -7,9 +7,6 @@ from .handler import SkipHandler log = logging.getLogger('aiogram.StateMachine') -# TODO: Provide async storage - - class BaseStorage: """ Skeleton for states storage @@ -132,6 +129,115 @@ class BaseStorage: raise NotImplementedError +class BaseAsyncStorage(BaseStorage): + async def set_state(self, chat, user, state): + """ + Set state + + :param chat: chat_id + :param user: user_id + :param state: value + """ + raise NotImplementedError + + async def get_state(self, chat, user): + """ + Get user state from + + :param chat: + :param user: + :return: + """ + raise NotImplementedError + + async def del_state(self, chat, user): + """ + Clear user state + :param chat: cha + :param user: + :return: + """ + raise NotImplementedError + + async def all_states(self, chat=None, user=None, state=None): + """ + Yield all states (Can use filters) + + :param chat: + :param user: + :param state: + :return: + """ + raise NotImplementedError + + async def set_value(self, chat, user, key, value): + """ + Set value for user in storage + + :param chat: + :param user: + :param key: + :param value: + :return: + """ + raise NotImplementedError + + async def get_value(self, chat, user, key, default=None): + """ + Get value from storage + + By default, this method calls `(await self.get_data(chat, user)).get(key, default)` + :param chat: + :param user: + :param key: + :param default: + :return: + """ + return (await self.get_data(chat, user)).get(key, default) + + async def del_value(self, chat, user, key): + """ + Delete value from storage + + :param chat: + :param user: + :param key: + """ + raise NotImplementedError + + async def get_data(self, chat, user): + """ + Get all stored data for user + + :param chat: + :param user: + :return: dict + """ + raise NotImplementedError + + async def update_data(self, chat, user, data): + """ + Update data in storage + + :param chat: + :param user: + :param data: + :return: + """ + raise NotImplementedError + + async def clear_data(self, chat, user, key): + """ + Clear data in storage + + :param chat: + :param user: + :param key: + :return: + """ + raise NotImplementedError + + class MemoryStorage(BaseStorage): """ Simple in-memory state storage @@ -398,6 +504,116 @@ class Controller: return f"{self._chat}:{self._user} - {self._state}" +class AsyncController: + """ + Storage controller + + Make easy access from callback's + """ + + def __init__(self, state_machine, chat, user, state): + self._state_machine = state_machine + self._chat = chat + self._user = user + self._state = state + + async def set_state(self, value): + """ + Set state + + :param value: + :return: + """ + await self._state_machine.set_state(self._chat, self._user, value) + + async def get_state(self): + """ + Get current state + + :return: + """ + return await self._state_machine.get_state(self._chat, self._user) + + async def clear(self): + """ + Reset state + + :return: + """ + await self._state_machine.del_state(self._chat, self._user) + + async def get(self, key, default=None): + """ + Get value from storage + + :param key: + :param default: + :return: + """ + return await self._state_machine.storage.get_value(self._chat, self._user, key, default) + + async def pop(self, key, default=None): + """ + Pop item from storage + + :param key: + :param default: + :return: + """ + result = await self.get(key, default) + await self.delete(key) + return result + + async def set(self, key, value): + """ + Set new value in user storage + + :param key: + :param value: + :return: + """ + await self._state_machine.storage.set_value(self._chat, self._user, key, value) + + async def delete(self, key): + """ + Delete key from user storage + + :param key: + :return: + """ + await self._state_machine.storage.del_value(self._chat, self._user, key) + + async def update(self, data): + """ + Update user storage + + :param data: + :return: + """ + await self._state_machine.storage.update_data(self._chat, self._user, data) + + @property + async def data(self): + """ + User data + + :return: + """ + return await self._state_machine.storage.get_data(self._chat, self._user) + + def __setitem__(self, key, value): + raise RuntimeError("Item assignment not allowed with async storage") + + def __getitem__(self, item): + raise RuntimeError("Item assignment not allowed with async storage") + + def __delitem__(self, key): + raise RuntimeError("Item assignment not allowed with async storage") + + def __str__(self): + return f"{self._chat}:{self._user} - {self._state}" + + class StateMachine: """ Manage state @@ -478,3 +694,84 @@ class StateMachine: callback = self.steps[state] controller = Controller(self, chat_id, from_user_id, state) await callback(message, controller) + + +class AsyncStateMachine: + """ + Manage state + """ + + def __init__(self, dispatcher, states, storage=None): + assert isinstance(storage, BaseAsyncStorage) + + self.steps = self._prepare_states(states) + self.storage = storage + + dispatcher.message_handlers.register(self.process_message, index=0) + + @staticmethod + def _prepare_states(states): + if isinstance(states, dict): + return states + elif isinstance(states, (list, tuple, set)): + prepared_states = {} + for state in states: + if not callable(state): + raise TypeError('State must be an callable') + state_name = state.__name__ + prepared_states[state_name] = state + return prepared_states + raise TypeError('States must be an dict or list!') + + async def set_state(self, chat, user, state): + """ + Save state to storage + :param chat: + :param user: + :param state: + :return: + """ + log.debug(f"Set state for {chat}:{user} to '{state}'") + await self.storage.set_state(chat, user, state) + + async def get_state(self, chat, user): + """ + Get state from storage + :param chat: + :param user: + :return: + """ + return await self.storage.get_state(chat, user) + + async def del_state(self, chat, user): + """ + Clear user state + :param chat: + :param user: + :return: + """ + log.debug(f"Reset state for {chat}:{user}") + await self.storage.del_state(chat, user) + + async def process_message(self, message): + """ + Read message and process it + :param message: + :return: + """ + chat_id = message.chat.id + from_user_id = message.from_user.id + + state = await self.get_state(chat_id, from_user_id) + if state is None: + raise SkipHandler() + + if state not in self.steps: + log.warning(f"Found unknown state '{state}' for {chat_id}:{from_user_id}. Condition will be reset.") + await self.del_state(chat_id, from_user_id) + raise SkipHandler() + + log.debug(f"Process state for {chat_id}:{from_user_id} - '{state}'") + callback = self.steps[state] + controller = AsyncController(self, chat_id, from_user_id, state) + await callback(message, controller)