diff --git a/pyrogram/dispatcher.py b/pyrogram/dispatcher.py index 6e503cebb3..77e21e77d8 100644 --- a/pyrogram/dispatcher.py +++ b/pyrogram/dispatcher.py @@ -26,7 +26,8 @@ from pyrogram.handlers import ( CallbackQueryHandler, MessageHandler, EditedMessageHandler, DeletedMessagesHandler, UserStatusHandler, RawUpdateHandler, InlineQueryHandler, PollHandler, - ChosenInlineResultHandler, ChatMemberUpdatedHandler, ChatJoinRequestHandler + ChosenInlineResultHandler, ChatMemberUpdatedHandler, ChatJoinRequestHandler, + ErrorHandler ) from pyrogram.raw.types import ( UpdateNewMessage, UpdateNewChannelMessage, UpdateNewScheduledMessage, @@ -62,6 +63,7 @@ def __init__(self, client: "pyrogram.Client"): self.updates_queue = asyncio.Queue() self.groups = OrderedDict() + self.error_handlers = [] async def message_parser(update, users, chats): return ( @@ -163,6 +165,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) @@ -172,11 +175,15 @@ 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() @@ -189,10 +196,16 @@ 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.") + 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) + self.groups[group].remove(handler) finally: for lock in self.locks_list: lock.release() @@ -250,8 +263,32 @@ async def handler_worker(self, lock): except pyrogram.ContinuePropagation: continue except Exception as e: - log.exception(e) - + handled_error = False + for error_handler in self.error_handlers: + try: + if await error_handler.check(self.client, e): + if inspect.iscoroutinefunction(error_handler.callback): + await error_handler.callback(self.client, e, *args) + else: + await self.loop.run_in_executor( + self.client.executor, + error_handler.callback, + self.client, + e, + *args + ) + handled_error = True + break + except pyrogram.StopPropagation: + raise + except pyrogram.ContinuePropagation: + continue + except Exception as e: + log.exception(e) + continue + + if not handled_error: + log.exception(e) break except pyrogram.StopPropagation: pass diff --git a/pyrogram/handlers/__init__.py b/pyrogram/handlers/__init__.py index 1c762958a0..67a61f9a85 100644 --- a/pyrogram/handlers/__init__.py +++ b/pyrogram/handlers/__init__.py @@ -28,3 +28,4 @@ from .poll_handler import PollHandler from .raw_update_handler import RawUpdateHandler from .user_status_handler import UserStatusHandler +from .error_handler import ErrorHandler diff --git a/pyrogram/handlers/error_handler.py b/pyrogram/handlers/error_handler.py new file mode 100644 index 0000000000..307002545e --- /dev/null +++ b/pyrogram/handlers/error_handler.py @@ -0,0 +1,64 @@ +# Pyrogram - Telegram MTProto API Client Library for Python +# Copyright (C) 2017-present Dan +# +# This file is part of Pyrogram. +# +# Pyrogram 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. +# +# Pyrogram 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 Pyrogram. If not, see . +from typing import Callable +import pyrogram +from .handler import Handler + + +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_message` 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). + + errors (:obj:`Exception` | List of :obj:`Exception`): + 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 message handler. + + error (:obj:`~Exception`): + The error that was raised. + + update (:obj:`~pyrogram.Update`): + The update that caused the error. + """ + + def __init__(self, callback: Callable, errors=None): + if errors is None: + errors = [Exception] + else: + if not isinstance(errors, list): + errors = [errors] + + self.errors = errors + super().__init__(callback) + + async def check(self, client: "pyrogram.Client", error: Exception): + return any(isinstance(error, e) for e in self.errors) + + def check_remove(self, error: Exception): + return self.errors == error or any(isinstance(error, e) for e in self.errors) diff --git a/pyrogram/methods/decorators/__init__.py b/pyrogram/methods/decorators/__init__.py index 1fc61185c9..02f54592be 100644 --- a/pyrogram/methods/decorators/__init__.py +++ b/pyrogram/methods/decorators/__init__.py @@ -28,6 +28,7 @@ from .on_poll import OnPoll from .on_raw_update import OnRawUpdate from .on_user_status import OnUserStatus +from .on_error import OnError class Decorators( @@ -42,6 +43,7 @@ class Decorators( OnPoll, OnChosenInlineResult, OnChatMemberUpdated, - OnChatJoinRequest + OnChatJoinRequest, + OnError ): pass diff --git a/pyrogram/methods/decorators/on_error.py b/pyrogram/methods/decorators/on_error.py new file mode 100644 index 0000000000..aa7e15d713 --- /dev/null +++ b/pyrogram/methods/decorators/on_error.py @@ -0,0 +1,49 @@ +# Pyrogram - Telegram MTProto API Client Library for Python +# Copyright (C) 2017-present Dan +# +# This file is part of Pyrogram. +# +# Pyrogram 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. +# +# Pyrogram 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 Pyrogram. 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 80a5f7419d..7743cbfcad 100644 --- a/pyrogram/methods/utilities/__init__.py +++ b/pyrogram/methods/utilities/__init__.py @@ -24,16 +24,18 @@ from .start import Start from .stop import Stop from .stop_transmission import StopTransmission +from .remove_error_handler import RemoveErrorHandler class Utilities( AddHandler, ExportSessionString, RemoveHandler, + RemoveErrorHandler, Restart, Run, Start, Stop, - StopTransmission + StopTransmission, ): pass diff --git a/pyrogram/methods/utilities/remove_error_handler.py b/pyrogram/methods/utilities/remove_error_handler.py new file mode 100644 index 0000000000..787e31ad24 --- /dev/null +++ b/pyrogram/methods/utilities/remove_error_handler.py @@ -0,0 +1,35 @@ +# Pyrogram - Telegram MTProto API Client Library for Python +# Copyright (C) 2017-present Dan +# +# This file is part of Pyrogram. +# +# Pyrogram 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. +# +# Pyrogram 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 Pyrogram. If not, see . + +import pyrogram + + +class RemoveErrorHandler: + def remove_error_handler( + self: "pyrogram.Client", + error: "Exception | tuple[Exception]" = Exception + ): + """Remove a previously-registered error handler. (using exception classes) + + Parameters: + error (``Exception``): + The error(s) for handlers to be removed. + """ + for handler in self.dispatcher.error_handlers: + if handler.check_remove(error): + self.dispatcher.error_handlers.remove(handler)