fix: return default state on get_state

This commit is contained in:
Oleg A 2021-03-25 15:41:20 +03:00
parent 5d851125f7
commit 1dfc00ee4d
4 changed files with 7 additions and 5 deletions

View file

@ -35,7 +35,7 @@ class MemoryStorage(BaseStorage):
user: typing.Union[str, int, None] = None,
default: typing.Optional[str] = None) -> typing.Optional[str]:
chat, user = self.resolve_address(chat=chat, user=user)
return self.data[chat][user]['state']
return self.data[chat][user].get("state", self.resolve_state(default))
async def get_data(self, *,
chat: typing.Union[str, int, None] = None,

View file

@ -136,7 +136,7 @@ class MongoStorage(BaseStorage):
db = await self.get_db()
result = await db[STATE].find_one(filter={'chat': chat, 'user': user})
return result.get('state') if result else default
return result.get('state') if result else self.resolve_state(default)
async def set_data(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None,
data: Dict = None):

View file

@ -118,7 +118,7 @@ class RedisStorage(BaseStorage):
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]:
record = await self.get_record(chat=chat, user=user)
return record['state']
return record.get('state', self.resolve_state(default))
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:
@ -277,7 +277,7 @@ class RedisStorage2(BaseStorage):
chat, user = self.check_address(chat=chat, user=user)
key = self.generate_key(chat, user, STATE_KEY)
redis = await self.redis()
return await redis.get(key, encoding='utf8') or None
return await redis.get(key, encoding='utf8') or self.resolve_state(default)
async def get_data(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None,
default: typing.Optional[dict] = None) -> typing.Dict:

View file

@ -95,7 +95,9 @@ class RethinkDBStorage(BaseStorage):
default: typing.Optional[str] = None) -> typing.Optional[str]:
chat, user = map(str, self.check_address(chat=chat, user=user))
async with self.connection() as conn:
return await r.table(self._table).get(chat)[user]['state'].default(default or None).run(conn)
return await r.table(self._table).get(chat)[user]['state'].default(
self.resolve_state(default) or None
).run(conn)
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: