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

Add decorators, return int. names in mongo

Browse source
This commit is contained in:
Anthony Byuraev 2020-07-05 15:47:23 +03:00
parent 352588cb06
commit 860ad4c169
1 changed files with 58 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

84
aiogram/contrib/fsm_storage/mongo.py
View file

@ -16,6 +16,7 @@ except ModuleNotFoundError as e:
raise e
from ...dispatcher.storage import BaseStorage
from aiogram.utils.deprecated import renamed_argument
STATE = 'aiogram_state'
DATA = 'aiogram_data'
@ -114,75 +115,106 @@ 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,
@renamed_argument(old_name='user', new_name='user_id', until_version='3.0', stacklevel=3)
@renamed_argument(old_name='chat', new_name='chat_id', until_version='3.0', stacklevel=4)
async def set_state(self, *,
chat_id: Union[str, int, None] = None,
user_id: Union[str, int, None] = None,
state: Optional[AnyStr] = None):
chat, user = self.check_address(chat=chat, user=user)
chat_id, user_id = self.check_address(chat_id=chat_id, user_id=user_id)
db = await self.get_db()
if state is None:
await db[STATE].delete_one(filter={'chat': chat, 'user': user})
await db[STATE].delete_one(filter={'chat': chat_id, 'user': user_id})
else:
await db[STATE].update_one(filter={'chat': chat, 'user': user},
await db[STATE].update_one(filter={'chat': chat_id, 'user': user_id},
update={'$set': {'state': state}}, upsert=True)
async def get_state(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None,
@renamed_argument(old_name='user', new_name='user_id', until_version='3.0', stacklevel=3)
@renamed_argument(old_name='chat', new_name='chat_id', until_version='3.0', stacklevel=4)
async def get_state(self, *,
chat_id: Union[str, int, None] = None,
user_id: Union[str, int, None] = None,
default: Optional[str] = None) -> Optional[str]:
chat, user = self.check_address(chat=chat, user=user)
chat_id, user_id = self.check_address(chat_id=chat_id, user_id=user_id)
db = await self.get_db()
result = await db[STATE].find_one(filter={'chat': chat, 'user': user})
result = await db[STATE].find_one(filter={'chat': chat_id, 'user': user_id})
return result.get('state') if result else default
async def set_data(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None,
@renamed_argument(old_name='user', new_name='user_id', until_version='3.0', stacklevel=3)
@renamed_argument(old_name='chat', new_name='chat_id', until_version='3.0', stacklevel=4)
async def set_data(self, *,
chat_id: Union[str, int, None] = None,
user_id: Union[str, int, None] = None,
data: Dict = None):
chat, user = self.check_address(chat=chat, user=user)
chat_id, user_id = self.check_address(chat_id=chat_id, user_id=user_id)
db = await self.get_db()
await db[DATA].update_one(filter={'chat': chat, 'user': user},
await db[DATA].update_one(filter={'chat': chat_id, 'user': user_id},
update={'$set': {'data': data}}, upsert=True)
async def get_data(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None,
@renamed_argument(old_name='user', new_name='user_id', until_version='3.0', stacklevel=3)
@renamed_argument(old_name='chat', new_name='chat_id', until_version='3.0', stacklevel=4)
async def get_data(self, *,
chat_id: Union[str, int, None] = None,
user_id: Union[str, int, None] = None,
default: Optional[dict] = None) -> Dict:
chat, user = self.check_address(chat=chat, user=user)
chat_id, user_id = self.check_address(chat_id=chat_id, user_id=user_id)
db = await self.get_db()
result = await db[DATA].find_one(filter={'chat': chat, 'user': user})
result = await db[DATA].find_one(filter={'chat': chat_id, 'user': user_id})
return result.get('data') if result else default or {}
async def update_data(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None,
@renamed_argument(old_name='user', new_name='user_id', until_version='3.0', stacklevel=3)
@renamed_argument(old_name='chat', new_name='chat_id', until_version='3.0', stacklevel=4)
async def update_data(self, *,
chat_id: Union[str, int, None] = None,
user_id: Union[str, int, None] = None,
data: Dict = None, **kwargs):
if data is None:
data = {}
temp_data = await self.get_data(chat=chat, user=user, default={})
temp_data = await self.get_data(chat_id=chat_id, user_id=user_id, default={})
temp_data.update(data, **kwargs)
await self.set_data(chat=chat, user=user, data=temp_data)
await self.set_data(chat_id=chat_id, user_id=user_id, data=temp_data)
def has_bucket(self):
return True
async def get_bucket(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None,
@renamed_argument(old_name='user', new_name='user_id', until_version='3.0', stacklevel=3)
@renamed_argument(old_name='chat', new_name='chat_id', until_version='3.0', stacklevel=4)
async def get_bucket(self, *,
chat_id: Union[str, int, None] = None,
user_id: Union[str, int, None] = None,
default: Optional[dict] = None) -> Dict:
chat, user = self.check_address(chat=chat, user=user)
chat_id, user_id = self.check_address(chat_id=chat_id, user_id=user_id)
db = await self.get_db()
result = await db[BUCKET].find_one(filter={'chat': chat, 'user': user})
result = await db[BUCKET].find_one(filter={'chat': chat_id, 'user': user_id})
return result.get('bucket') if result else default or {}
async def set_bucket(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None,
@renamed_argument(old_name='user', new_name='user_id', until_version='3.0', stacklevel=3)
@renamed_argument(old_name='chat', new_name='chat_id', until_version='3.0', stacklevel=4)
async def set_bucket(self, *,
chat_id: Union[str, int, None] = None,
user_id: Union[str, int, None] = None,
bucket: Dict = None):
chat, user = self.check_address(chat=chat, user=user)
chat_id, user_id = self.check_address(chat_id=chat_id, user_id=user_id)
db = await self.get_db()
await db[BUCKET].update_one(filter={'chat': chat, 'user': user},
await db[BUCKET].update_one(filter={'chat': chat_id, 'user': user_id},
update={'$set': {'bucket': bucket}}, upsert=True)
async def update_bucket(self, *, chat: Union[str, int, None] = None,
user: Union[str, int, None] = None,
@renamed_argument(old_name='user', new_name='user_id', until_version='3.0', stacklevel=3)
@renamed_argument(old_name='chat', new_name='chat_id', until_version='3.0', stacklevel=4)
async def update_bucket(self, *,
chat_id: Union[str, int, None] = None,
user_id: Union[str, int, None] = None,
bucket: Dict = None, **kwargs):
if bucket is None:
bucket = {}
temp_bucket = await self.get_bucket(chat=chat, user=user)
temp_bucket = await self.get_bucket(chat_id=chat_id, user_id=user_id)
temp_bucket.update(bucket, **kwargs)
await self.set_bucket(chat=chat, user=user, bucket=temp_bucket)
await self.set_bucket(chat_id=chat_id, user_id=user_id, bucket=temp_bucket)
async def reset_all(self, full=True):
"""