From 5d851125f768707d48485948e2223bc5e743645b Mon Sep 17 00:00:00 2001 From: Oleg A Date: Thu, 25 Mar 2021 15:35:23 +0300 Subject: [PATCH] refactor: moved state resolution to storage --- aiogram/contrib/fsm_storage/memory.py | 4 ++-- aiogram/contrib/fsm_storage/mongo.py | 4 ++-- aiogram/contrib/fsm_storage/redis.py | 6 +++--- aiogram/contrib/fsm_storage/rethinkdb.py | 4 ++-- aiogram/dispatcher/storage.py | 22 +++++++++++----------- 5 files changed, 20 insertions(+), 20 deletions(-) diff --git a/aiogram/contrib/fsm_storage/memory.py b/aiogram/contrib/fsm_storage/memory.py index cd97f012..9343c6f3 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, 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, diff --git a/aiogram/contrib/fsm_storage/mongo.py b/aiogram/contrib/fsm_storage/mongo.py index c9aff88a..ab43963f 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, 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, ) diff --git a/aiogram/contrib/fsm_storage/redis.py b/aiogram/contrib/fsm_storage/redis.py index 28ab0bae..b32dacaa 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, 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): diff --git a/aiogram/contrib/fsm_storage/rethinkdb.py b/aiogram/contrib/fsm_storage/rethinkdb.py index c1a4f47a..6d7e2109 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, 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) diff --git a/aiogram/dispatcher/storage.py b/aiogram/dispatcher/storage.py index 5a7d101e..eb248e34 100644 --- a/aiogram/dispatcher/storage.py +++ b/aiogram/dispatcher/storage.py @@ -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)