From d5a09b06e18949158cd3a4024f7d78dcb125ecc9 Mon Sep 17 00:00:00 2001 From: Oleg A Date: Thu, 25 Mar 2021 15:19:15 +0300 Subject: [PATCH] fix: resolve `filters.State` on `set_state` --- aiogram/contrib/fsm_storage/memory.py | 4 ++-- aiogram/contrib/fsm_storage/mongo.py | 15 ++++++++++----- aiogram/contrib/fsm_storage/redis.py | 7 +++++-- aiogram/contrib/fsm_storage/rethinkdb.py | 11 ++++++++--- 4 files changed, 25 insertions(+), 12 deletions(-) diff --git a/aiogram/contrib/fsm_storage/memory.py b/aiogram/contrib/fsm_storage/memory.py index 2940f3fa..cd97f012 100644 --- a/aiogram/contrib/fsm_storage/memory.py +++ b/aiogram/contrib/fsm_storage/memory.py @@ -1,7 +1,7 @@ import copy import typing -from ...dispatcher.storage import BaseStorage +from ...dispatcher.storage import BaseStorage, FSMContext class MemoryStorage(BaseStorage): @@ -58,7 +58,7 @@ class MemoryStorage(BaseStorage): user: typing.Union[str, int, None] = None, state: typing.AnyStr = None): chat, user = self.resolve_address(chat=chat, user=user) - self.data[chat][user]['state'] = state + self.data[chat][user]['state'] = FSMContext.resolve_state(state) async def set_data(self, *, chat: typing.Union[str, int, None] = None, diff --git a/aiogram/contrib/fsm_storage/mongo.py b/aiogram/contrib/fsm_storage/mongo.py index f810a3eb..c9aff88a 100644 --- a/aiogram/contrib/fsm_storage/mongo.py +++ b/aiogram/contrib/fsm_storage/mongo.py @@ -15,7 +15,7 @@ except ModuleNotFoundError as e: warnings.warn("Install motor with `pip install motor`") raise e -from ...dispatcher.storage import BaseStorage +from ...dispatcher.storage import BaseStorage, FSMContext STATE = 'aiogram_state' DATA = 'aiogram_data' @@ -65,7 +65,7 @@ class MongoStorage(BaseStorage): try: self._mongo = AsyncIOMotorClient(self._uri) except pymongo.errors.ConfigurationError as e: - if "query() got an unexpected keyword argument 'lifetime'" in e.args[0]: + if "query() got an unexpected keyword argument 'lifetime'" in e.args[0]: import logging logger = logging.getLogger("aiogram") logger.warning("Run `pip install dnspython==1.16.0` in order to fix ConfigurationError. More information: https://github.com/mongodb/mongo-python-driver/pull/423#issuecomment-528998245") @@ -114,7 +114,9 @@ class MongoStorage(BaseStorage): async def wait_closed(self): return True - async def set_state(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None, + async def set_state(self, *, + chat: Union[str, int, None] = None, + user: Union[str, int, None] = None, state: Optional[AnyStr] = None): chat, user = self.check_address(chat=chat, user=user) db = await self.get_db() @@ -122,8 +124,11 @@ class MongoStorage(BaseStorage): if state is None: await db[STATE].delete_one(filter={'chat': chat, 'user': user}) else: - await db[STATE].update_one(filter={'chat': chat, 'user': user}, - update={'$set': {'state': state}}, upsert=True) + await db[STATE].update_one( + filter={'chat': chat, 'user': user}, + update={'$set': {'state': FSMContext.resolve_state(state)}}, + upsert=True, + ) async def get_state(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None, default: Optional[str] = None) -> Optional[str]: diff --git a/aiogram/contrib/fsm_storage/redis.py b/aiogram/contrib/fsm_storage/redis.py index 74dd736c..28ab0bae 100644 --- a/aiogram/contrib/fsm_storage/redis.py +++ b/aiogram/contrib/fsm_storage/redis.py @@ -8,7 +8,7 @@ import typing import aioredis -from ...dispatcher.storage import BaseStorage +from ...dispatcher.storage import BaseStorage, FSMContext from ...utils import json STATE_KEY = 'state' @@ -125,9 +125,12 @@ class RedisStorage(BaseStorage): record = await self.get_record(chat=chat, user=user) return record['data'] - async def set_state(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, + async def set_state(self, *, + chat: typing.Union[str, int, None] = None, + user: typing.Union[str, int, None] = None, state: typing.Optional[typing.AnyStr] = None): record = await self.get_record(chat=chat, user=user) + state = FSMContext.resolve_state(state) await self.set_record(chat=chat, user=user, state=state, data=record['data']) async def set_data(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, diff --git a/aiogram/contrib/fsm_storage/rethinkdb.py b/aiogram/contrib/fsm_storage/rethinkdb.py index 5bb9062a..c1a4f47a 100644 --- a/aiogram/contrib/fsm_storage/rethinkdb.py +++ b/aiogram/contrib/fsm_storage/rethinkdb.py @@ -5,7 +5,7 @@ import typing import rethinkdb from rethinkdb.asyncio_net.net_asyncio import Connection -from ...dispatcher.storage import BaseStorage +from ...dispatcher.storage import BaseStorage, FSMContext __all__ = ('RethinkDBStorage',) @@ -103,11 +103,16 @@ class RethinkDBStorage(BaseStorage): async with self.connection() as conn: return await r.table(self._table).get(chat)[user]['data'].default(default or {}).run(conn) - async def set_state(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, + async def set_state(self, *, + chat: typing.Union[str, int, None] = None, + user: typing.Union[str, int, None] = None, state: typing.Optional[typing.AnyStr] = None): chat, user = map(str, self.check_address(chat=chat, user=user)) async with self.connection() as conn: - await r.table(self._table).insert({'id': chat, user: {'state': state}}, conflict="update").run(conn) + await r.table(self._table).insert( + {'id': chat, user: {'state': FSMContext.resolve_state(state)}}, + conflict="update", + ).run(conn) async def set_data(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None, data: typing.Dict = None):