python/aiogram
1
0
Fork 0
mirror of https://github.com/aiogram/aiogram.git synced 2026-04-08 16:37:47 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

refactor: moved state resolution to storage

Browse source
This commit is contained in:
Oleg A 2021-03-25 15:35:23 +03:00
parent d5a09b06e1
commit 5d851125f7
5 changed files with 20 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

4
aiogram/contrib/fsm_storage/memory.py
View file

@ -1,7 +1,7 @@
import copy
import typing
from ...dispatcher.storage import BaseStorage, FSMContext
from ...dispatcher.storage import BaseStorage
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'] = FSMContext.resolve_state(state)
self.data[chat][user]['state'] = self.resolve_state(state)
async def set_data(self, *,
chat: typing.Union[str, int, None] = None,

4
aiogram/contrib/fsm_storage/mongo.py
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, FSMContext
from ...dispatcher.storage import BaseStorage
STATE = 'aiogram_state'
DATA = 'aiogram_data'
@ -126,7 +126,7 @@ class MongoStorage(BaseStorage):
else:
await db[STATE].update_one(
filter={'chat': chat, 'user': user},
update={'$set': {'state': FSMContext.resolve_state(state)}},
update={'$set': {'state': self.resolve_state(state)}},
upsert=True,
)

6
aiogram/contrib/fsm_storage/redis.py
View file

@ -8,7 +8,7 @@ import typing
import aioredis
from ...dispatcher.storage import BaseStorage, FSMContext
from ...dispatcher.storage import BaseStorage
from ...utils import json
STATE_KEY = 'state'
@ -130,7 +130,7 @@ class RedisStorage(BaseStorage):
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)
state = self.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,
@ -297,7 +297,7 @@ class RedisStorage2(BaseStorage):
if state is None:
await redis.delete(key)
else:
await redis.set(key, state, expire=self._state_ttl)
await redis.set(key, self.resolve_state(state), expire=self._state_ttl)
async def set_data(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None,
data: typing.Dict = None):

4
aiogram/contrib/fsm_storage/rethinkdb.py
View file

@ -5,7 +5,7 @@ import typing
import rethinkdb
from rethinkdb.asyncio_net.net_asyncio import Connection
from ...dispatcher.storage import BaseStorage, FSMContext
from ...dispatcher.storage import BaseStorage
__all__ = ('RethinkDBStorage',)
@ -110,7 +110,7 @@ class RethinkDBStorage(BaseStorage):
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': FSMContext.resolve_state(state)}},
{'id': chat, user: {'state': self.resolve_state(state)}},
conflict="update",
).run(conn)

22
aiogram/dispatcher/storage.py
View file

@ -273,15 +273,6 @@ class BaseStorage:
"""
await self.set_data(chat=chat, user=user, data={})
class FSMContext:
def __init__(self, storage, chat, user):
self.storage: BaseStorage = storage
self.chat, self.user = self.storage.check_address(chat=chat, user=user)
def proxy(self):
return FSMContextProxy(self)
@staticmethod
def resolve_state(value):
from .filters.state import State
@ -297,8 +288,17 @@ class FSMContext:
return str(value)
class FSMContext:
def __init__(self, storage, chat, user):
self.storage: BaseStorage = storage
self.chat, self.user = self.storage.check_address(chat=chat, user=user)
def proxy(self):
return FSMContextProxy(self)
async def get_state(self, default: typing.Optional[str] = None) -> typing.Optional[str]:
return await self.storage.get_state(chat=self.chat, user=self.user, default=self.resolve_state(default))
return await self.storage.get_state(chat=self.chat, user=self.user, default=default)
async def get_data(self, default: typing.Optional[str] = None) -> typing.Dict:
return await self.storage.get_data(chat=self.chat, user=self.user, default=default)
@ -307,7 +307,7 @@ class FSMContext:
await self.storage.update_data(chat=self.chat, user=self.user, data=data, **kwargs)
async def set_state(self, state: typing.Optional[typing.AnyStr] = None):
await self.storage.set_state(chat=self.chat, user=self.user, state=self.resolve_state(state))
await self.storage.set_state(chat=self.chat, user=self.user, state=state)
async def set_data(self, data: typing.Dict = None):
await self.storage.set_data(chat=self.chat, user=self.user, data=data)