From 86ad9ea6162f10e3fd0246639898268293cb857d Mon Sep 17 00:00:00 2001 From: Yasir Aris M Date: Fri, 20 Dec 2024 19:59:22 +0700 Subject: [PATCH] pyrofork: Add Error handler Signed-off-by: Yasir Aris M --- docs/source/api/decorators.rst | 1 + pyrogram/dispatcher.py | 159 ++++++++++-------- pyrogram/handlers/__init__.py | 1 + pyrogram/handlers/error_handler.py | 79 +++++++++ pyrogram/methods/decorators/__init__.py | 2 + pyrogram/methods/decorators/on_error.py | 50 ++++++ pyrogram/methods/utilities/__init__.py | 2 + .../methods/utilities/remove_error_handler.py | 42 +++++ 8 files changed, 268 insertions(+), 68 deletions(-) create mode 100644 pyrogram/handlers/error_handler.py create mode 100644 pyrogram/methods/decorators/on_error.py create mode 100644 pyrogram/methods/utilities/remove_error_handler.py diff --git a/docs/source/api/decorators.rst b/docs/source/api/decorators.rst index d08f6cb52..2c326ed9a 100644 --- a/docs/source/api/decorators.rst +++ b/docs/source/api/decorators.rst @@ -82,3 +82,4 @@ Details .. autodecorator:: pyrogram.Client.on_poll() .. autodecorator:: pyrogram.Client.on_disconnect() .. autodecorator:: pyrogram.Client.on_raw_update() +.. autodecorator:: pyrogram.Client.on_error() diff --git a/pyrogram/dispatcher.py b/pyrogram/dispatcher.py index a0ee6cb93..592f19524 100644 --- a/pyrogram/dispatcher.py +++ b/pyrogram/dispatcher.py @@ -23,9 +23,8 @@ from collections import OrderedDict import pyrogram -from pyrogram import errors -from pyrogram import utils -from pyrogram import raw +from pyrogram import errors, raw, types, utils +from pyrogram.handlers.handler import Handler from pyrogram.handlers import ( BotBusinessConnectHandler, BotBusinessMessageHandler, @@ -33,6 +32,7 @@ MessageHandler, EditedMessageHandler, EditedBotBusinessMessageHandler, + ErrorHandler, DeletedMessagesHandler, DeletedBotBusinessMessagesHandler, MessageReactionUpdatedHandler, @@ -97,6 +97,7 @@ def __init__(self, client: "pyrogram.Client"): self.handler_worker_tasks = [] self.locks_list = [] + self.error_handlers = [] self.updates_queue = asyncio.Queue() self.groups = OrderedDict() @@ -286,6 +287,7 @@ async def stop(self): self.handler_worker_tasks.clear() self.groups.clear() + self.error_handlers.clear() log.info("Stopped %s HandlerTasks", self.client.workers) @@ -295,11 +297,14 @@ async def fn(): await lock.acquire() try: - if group not in self.groups: - self.groups[group] = [] - self.groups = OrderedDict(sorted(self.groups.items())) - - self.groups[group].append(handler) + if isinstance(handler, ErrorHandler): + if handler not in self.error_handlers: + self.error_handlers.append(handler) + else: + if group not in self.groups: + self.groups[group] = [] + self.groups = OrderedDict(sorted(self.groups.items())) + self.groups[group].append(handler) finally: for lock in self.locks_list: lock.release() @@ -312,76 +317,94 @@ async def fn(): await lock.acquire() try: - if group not in self.groups: - raise ValueError(f"Group {group} does not exist. Handler was not removed.") - - self.groups[group].remove(handler) + if isinstance(handler, ErrorHandler): + if handler not in self.error_handlers: + raise ValueError( + f"Error handler {handler} does not exist. Handler was not removed." + ) + self.error_handlers.remove(handler) + else: + if group not in self.groups: + raise ValueError(f"Group {group} does not exist. Handler was not removed.") + self.groups[group].remove(handler) finally: for lock in self.locks_list: lock.release() self.loop.create_task(fn()) - async def handler_worker(self, lock): + async def handler_worker(self, lock: asyncio.Lock): while True: packet = await self.updates_queue.get() - if packet is None: break - - try: - update, users, chats = packet - parser = self.update_parsers.get(type(update), None) - - parsed_update, handler_type = ( - await parser(update, users, chats) - if parser is not None - else (None, type(None)) - ) - - async with lock: - for group in self.groups.values(): - for handler in group: - args = None - - if isinstance(handler, handler_type): - try: - if await handler.check(self.client, parsed_update): - args = (parsed_update,) - except Exception as e: - log.exception(e) - continue - + await self._process_packet(packet, lock) + + async def _process_packet( + self, + packet: tuple[raw.core.TLObject, dict[int, types.Update], dict[int, types.Update]], + lock: asyncio.Lock, + ): + try: + update, users, chats = packet + parser = self.update_parsers.get(type(update)) + + if parser is not None: + parsed_result = parser(update, users, chats) + if inspect.isawaitable(parsed_result): + parsed_update, handler_type = await parsed_result + else: + parsed_update, handler_type = parsed_result + else: + parsed_update, handler_type = (None, type(None)) + + async with lock: + for group in self.groups.values(): + for handler in group: + try: + if parsed_update is not None: + if isinstance(handler, handler_type) and await handler.check( + self.client, parsed_update + ): + await self._execute_callback(handler, parsed_update) + break elif isinstance(handler, RawUpdateHandler): - try: - if await handler.check(self.client, update): - args = (update, users, chats) - except Exception as e: - log.exception(e) - continue - - if args is None: - continue - - try: - if inspect.iscoroutinefunction(handler.callback): - await handler.callback(self.client, *args) - else: - await self.loop.run_in_executor( - self.client.executor, - handler.callback, - self.client, - *args - ) - except pyrogram.StopPropagation: + await self._execute_callback(handler, update, users, chats) + break + except (pyrogram.StopPropagation, pyrogram.ContinuePropagation) as e: + if isinstance(e, pyrogram.StopPropagation): raise - except pyrogram.ContinuePropagation: - continue - except Exception as e: - log.exception(e) - - break + except Exception as exception: + if parsed_update is not None: + await self._handle_exception(parsed_update, exception) + except pyrogram.StopPropagation: + pass + except Exception as e: + log.exception(e) + finally: + self.updates_queue.task_done() + + async def _handle_exception(self, parsed_update: types.Update, exception: Exception): + handled_error = False + for error_handler in self.error_handlers: + try: + if await error_handler.check(self.client, parsed_update, exception): + handled_error = True + break except pyrogram.StopPropagation: - pass - except Exception as e: - log.exception(e) + raise + except pyrogram.ContinuePropagation: + continue + except Exception as inner_exception: + log.exception("Error in error handler: %s", inner_exception) + + if not handled_error: + log.exception("Unhandled exception: %s", exception) + + async def _execute_callback(self, handler: Handler, *args): + if inspect.iscoroutinefunction(handler.callback): + await handler.callback(self.client, *args) + else: + await self.client.loop.run_in_executor( + self.client.executor, handler.callback, self.client, *args + ) diff --git a/pyrogram/handlers/__init__.py b/pyrogram/handlers/__init__.py index eb879a4d6..b4611cc6a 100644 --- a/pyrogram/handlers/__init__.py +++ b/pyrogram/handlers/__init__.py @@ -29,6 +29,7 @@ from .disconnect_handler import DisconnectHandler from .edited_message_handler import EditedMessageHandler from .edited_bot_business_message_handler import EditedBotBusinessMessageHandler +from .error_handler import ErrorHandler from .inline_query_handler import InlineQueryHandler from .message_handler import MessageHandler from .poll_handler import PollHandler diff --git a/pyrogram/handlers/error_handler.py b/pyrogram/handlers/error_handler.py new file mode 100644 index 000000000..cf45bc116 --- /dev/null +++ b/pyrogram/handlers/error_handler.py @@ -0,0 +1,79 @@ +# Pyrofork - Telegram MTProto API Client Library for Python +# Copyright (C) 2017-present Dan +# Copyright (C) 2022-present Mayuri-Chan +# +# This file is part of Pyrofork. +# +# Pyrofork is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# Pyrofork is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with Pyrofork. If not, see . + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Callable + +from .handler import Handler + +import pyrogram +from pyrogram.types import Update + + +class ErrorHandler(Handler): + """The Error handler class. Used to handle errors. + It is intended to be used with :meth:`~pyrogram.Client.add_handler` + + For a nicer way to register this handler, have a look at the + :meth:`~pyrogram.Client.on_error` decorator. + + Parameters: + callback (``Callable``): + Pass a function that will be called when a new Error arrives. It takes *(client, error)* + as positional arguments (look at the section below for a detailed description). + + exceptions (``Exception`` | Iterable of ``Exception``, *optional*): + Pass one or more exception classes to allow only a subset of errors to be passed + in your callback function. + + Other parameters: + client (:obj:`~pyrogram.Client`): + The Client itself, useful when you want to call other API methods inside the error handler. + + update (:obj:`~pyrogram.Update`): + The update that caused the error. + + error (``Exception``): + The error that was raised. + """ + + def __init__( + self, + callback: Callable, + exceptions: type[Exception] | Iterable[type[Exception]] | None = None, + ): + self.exceptions = ( + tuple(exceptions) + if isinstance(exceptions, Iterable) + else (exceptions,) + if exceptions + else (Exception,) + ) + super().__init__(callback) + + async def check(self, client: pyrogram.Client, update: Update, exception: Exception) -> bool: + if isinstance(exception, self.exceptions): + await self.callback(client, update, exception) + return True + return False + + def check_remove(self, error: type[Exception] | Iterable[type[Exception]]) -> bool: + return isinstance(error, self.exceptions) diff --git a/pyrogram/methods/decorators/__init__.py b/pyrogram/methods/decorators/__init__.py index 402b7e26d..e71207296 100644 --- a/pyrogram/methods/decorators/__init__.py +++ b/pyrogram/methods/decorators/__init__.py @@ -28,6 +28,7 @@ from .on_disconnect import OnDisconnect from .on_edited_message import OnEditedMessage from .on_edited_bot_business_message import OnEditedBotBusinessMessage +from .on_error import OnError from .on_inline_query import OnInlineQuery from .on_message import OnMessage from .on_poll import OnPoll @@ -47,6 +48,7 @@ class Decorators( OnBotBusinessMessage, OnEditedMessage, OnEditedBotBusinessMessage, + OnError, OnDeletedMessages, OnDeletedBotBusinessMessages, OnCallbackQuery, diff --git a/pyrogram/methods/decorators/on_error.py b/pyrogram/methods/decorators/on_error.py new file mode 100644 index 000000000..d21d75bb0 --- /dev/null +++ b/pyrogram/methods/decorators/on_error.py @@ -0,0 +1,50 @@ +# Pyrofork - Telegram MTProto API Client Library for Python +# Copyright (C) 2017-present Dan +# Copyright (C) 2022-present Mayuri-Chan +# +# This file is part of Pyrofork. +# +# Pyrofork is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# Pyrofork is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with Pyrofork. If not, see . + +from typing import Callable + +import pyrogram +from pyrogram.filters import Filter + + +class OnError: + def on_error(self=None, errors=None) -> Callable: + """Decorator for handling new errors. + + This does the same thing as :meth:`~pyrogram.Client.add_handler` using the + :obj:`~pyrogram.handlers.MessageHandler`. + + Parameters: + errors (:obj:`~Exception`, *optional*): + Pass one or more errors to allow only a subset of errors to be passed + in your function. + """ + + def decorator(func: Callable) -> Callable: + if isinstance(self, pyrogram.Client): + self.add_handler(pyrogram.handlers.ErrorHandler(func, errors), 0) + elif isinstance(self, Filter) or self is None: + if not hasattr(func, "handlers"): + func.handlers = [] + + func.handlers.append((pyrogram.handlers.ErrorHandler(func, self), 0)) + + return func + + return decorator diff --git a/pyrogram/methods/utilities/__init__.py b/pyrogram/methods/utilities/__init__.py index 2184025be..22d974342 100644 --- a/pyrogram/methods/utilities/__init__.py +++ b/pyrogram/methods/utilities/__init__.py @@ -20,6 +20,7 @@ from .add_handler import AddHandler from .export_session_string import ExportSessionString from .remove_handler import RemoveHandler +from .remove_error_handler import RemoveErrorHandler from .restart import Restart from .run import Run from .run_sync import RunSync @@ -32,6 +33,7 @@ class Utilities( AddHandler, ExportSessionString, RemoveHandler, + RemoveErrorHandler, Restart, Run, RunSync, diff --git a/pyrogram/methods/utilities/remove_error_handler.py b/pyrogram/methods/utilities/remove_error_handler.py new file mode 100644 index 000000000..7ad9f7b40 --- /dev/null +++ b/pyrogram/methods/utilities/remove_error_handler.py @@ -0,0 +1,42 @@ +# Pyrofork - Telegram MTProto API Client Library for Python +# Copyright (C) 2017-present Dan +# Copyright (C) 2022-present Mayuri-Chan +# +# This file is part of Pyrofork. +# +# Pyrofork is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# Pyrofork is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with Pyrofork. If not, see . + +from __future__ import annotations +from collections.abc import Iterable +import pyrogram + + +class RemoveErrorHandler: + def remove_error_handler( + self: pyrogram.Client, + exception: type[Exception] | Iterable[type[Exception]] = Exception, + ): + """Remove a previously registered error handler using exception classes. + + Parameters: + exception (``Exception`` | Iterable of ``Exception``, *optional*): + The error(s) for handlers to be removed. Defaults to Exception. + """ + to_remove = [ + handler + for handler in self.dispatcher.error_handlers + if handler.check_remove(exception) + ] + for handler in to_remove: + self.dispatcher.error_handlers.remove(handler)