outside-fetcher/tasks/executors/mtproto_task_abstraction.py

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)