fix: resolve filters.State on set_state

This commit is contained in:
Oleg A 2021-03-25 15:19:15 +03:00
parent 45405f979c
commit d5a09b06e1
4 changed files with 25 additions and 12 deletions

View file

@ -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,

View file

@ -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]:

View file

@ -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,

View file

@ -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):