import abc import typing from async_worker import AsyncTask from pyrogram.errors.exceptions import FloodWait, RPCError, InternalServerError class MtProtoTask(AsyncTask, abc.ABC): async def _execute(self, func: typing.Callable) -> typing.Any: try: return await func() except FloodWait as error: return int(error.MESSAGE.split("_")[-1]) * 1e9 except InternalServerError: return 1e9 except RPCError as error: return False if abs(error.CODE) < 500 else 2e9 async def _setup(self) -> bool: return await self._execute(super()._setup) async def _process(self) -> typing.Union[bool, int]: return await self._execute(super()._process)