python/aiogram
1
0
Fork 0
mirror of https://github.com/aiogram/aiogram.git synced 2025-12-11 01:54:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Update memory storage (FSM) for usage with throttling manager.

Browse source
This commit is contained in:
Alex Root Junior 2017-12-02 01:06:23 +02:00
parent 2ebd7c5de4
commit 435b498d25
1 changed files with 25 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -30,7 +30,7 @@ class MemoryStorage(BaseStorage):
chat_id = str(chat_id)
user_id = str(user_id)
if user_id not in self.data[chat_id]:
self.data[chat_id][user_id] = {'state': None, 'data': {}}
self.data[chat_id][user_id] = {'state': None, 'data': {}, 'bucket': {}}
return self.data[chat_id][user_id]
async def get_state(self, *,
@ -82,3 +82,27 @@ class MemoryStorage(BaseStorage):
await self.set_state(chat=chat, user=user, state=None)
if with_data:
await self.set_data(chat=chat, user=user, data={})
def has_bucket(self):
return True
async def get_bucket(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None,
default: typing.Optional[dict] = None) -> typing.Dict:
chat, user = self.check_address(chat=chat, user=user)
user = self._get_user(chat, user)
return user['bucket']
async def set_bucket(self, *, chat: typing.Union[str, int, None] = None, user: typing.Union[str, int, None] = None,
bucket: typing.Dict = None):
chat, user = self.check_address(chat=chat, user=user)
user = self._get_user(chat, user)
user['bucket'] = bucket
async def update_bucket(self, *, chat: typing.Union[str, int, None] = None,
user: typing.Union[str, int, None] = None,
bucket: typing.Dict = None, **kwargs):
chat, user = self.check_address(chat=chat, user=user)
user = self._get_user(chat, user)
if bucket is None:
bucket = []
user['bucket'].update(bucket, **kwargs)