From 82933849fb74fb2eeddf68a048e04d615c573f3b Mon Sep 17 00:00:00 2001 From: lordcodes Date: Tue, 21 Mar 2023 14:51:27 +0200 Subject: [PATCH 1/5] Error handlers --- pyrogram/dispatcher.py | 58 ++++++++++++++++++++----- pyrogram/handlers/__init__.py | 1 + pyrogram/handlers/error_handler.py | 58 +++++++++++++++++++++++++ pyrogram/methods/decorators/__init__.py | 4 +- pyrogram/methods/decorators/on_error.py | 49 +++++++++++++++++++++ 5 files changed, 158 insertions(+), 12 deletions(-) create mode 100644 pyrogram/handlers/error_handler.py create mode 100644 pyrogram/methods/decorators/on_error.py diff --git a/pyrogram/dispatcher.py b/pyrogram/dispatcher.py index 6e503cebb3..03ce42b2a8 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,31 @@ 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) + else: + await self.loop.run_in_executor( + self.client.executor, + error_handler.callback, + self.client, + e + ) + handled_error = True # If the error is handled, don't log it + 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..c3bcbcf13f --- /dev/null +++ b/pyrogram/handlers/error_handler.py @@ -0,0 +1,58 @@ +# 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. + """ + + 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) 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 From 1010b1e3741712ecfc6414b1d1e01b8aa9f27114 Mon Sep 17 00:00:00 2001 From: lordcodes Date: Tue, 21 Mar 2023 15:01:03 +0200 Subject: [PATCH 2/5] Error handlers update --- pyrogram/dispatcher.py | 7 ++++--- pyrogram/handlers/error_handler.py | 3 +++ 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/pyrogram/dispatcher.py b/pyrogram/dispatcher.py index 03ce42b2a8..77e21e77d8 100644 --- a/pyrogram/dispatcher.py +++ b/pyrogram/dispatcher.py @@ -268,15 +268,16 @@ async def handler_worker(self, lock): try: if await error_handler.check(self.client, e): if inspect.iscoroutinefunction(error_handler.callback): - await error_handler.callback(self.client, e) + await error_handler.callback(self.client, e, *args) else: await self.loop.run_in_executor( self.client.executor, error_handler.callback, self.client, - e + e, + *args ) - handled_error = True # If the error is handled, don't log it + handled_error = True break except pyrogram.StopPropagation: raise diff --git a/pyrogram/handlers/error_handler.py b/pyrogram/handlers/error_handler.py index c3bcbcf13f..e454ef5a4a 100644 --- a/pyrogram/handlers/error_handler.py +++ b/pyrogram/handlers/error_handler.py @@ -42,6 +42,9 @@ class ErrorHandler(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): From 38224ff413fcea3cc795a9f210e356aa1fc46b2f Mon Sep 17 00:00:00 2001 From: lordcodes Date: Tue, 21 Mar 2023 15:11:56 +0200 Subject: [PATCH 3/5] Remove error handler utility added --- pyrogram/handlers/error_handler.py | 3 ++ pyrogram/methods/utilities/__init__.py | 4 ++- .../methods/utilities/remove_error_handler.py | 35 +++++++++++++++++++ 3 files changed, 41 insertions(+), 1 deletion(-) create mode 100644 pyrogram/methods/utilities/remove_error_handler.py diff --git a/pyrogram/handlers/error_handler.py b/pyrogram/handlers/error_handler.py index e454ef5a4a..1c9f677898 100644 --- a/pyrogram/handlers/error_handler.py +++ b/pyrogram/handlers/error_handler.py @@ -59,3 +59,6 @@ def __init__(self, callback: Callable, errors=None): async def check(self, client: "pyrogram.Client", error: Exception): return any(isinstance(error, e) for e in self.errors) + + async 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/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..d85c1b0563 --- /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 await handler.check_remove(error): + self.dispatcher.error_handlers.remove(handler) From fcc8dea2c07de3333f853f0a3290958d958a0724 Mon Sep 17 00:00:00 2001 From: lordcodes <83734728+LORD-ME-CODE@users.noreply.github.com> Date: Wed, 22 Mar 2023 08:26:45 +0200 Subject: [PATCH 4/5] Update error_handler.py --- pyrogram/handlers/error_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrogram/handlers/error_handler.py b/pyrogram/handlers/error_handler.py index 1c9f677898..307002545e 100644 --- a/pyrogram/handlers/error_handler.py +++ b/pyrogram/handlers/error_handler.py @@ -60,5 +60,5 @@ def __init__(self, callback: Callable, errors=None): async def check(self, client: "pyrogram.Client", error: Exception): return any(isinstance(error, e) for e in self.errors) - async def check_remove(self, error: Exception): + def check_remove(self, error: Exception): return self.errors == error or any(isinstance(error, e) for e in self.errors) From e7ec3b8a497cd9dfa42f3654382f589340b64d83 Mon Sep 17 00:00:00 2001 From: lordcodes <83734728+LORD-ME-CODE@users.noreply.github.com> Date: Wed, 22 Mar 2023 08:27:44 +0200 Subject: [PATCH 5/5] Update remove_error_handler.py --- pyrogram/methods/utilities/remove_error_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrogram/methods/utilities/remove_error_handler.py b/pyrogram/methods/utilities/remove_error_handler.py index d85c1b0563..787e31ad24 100644 --- a/pyrogram/methods/utilities/remove_error_handler.py +++ b/pyrogram/methods/utilities/remove_error_handler.py @@ -31,5 +31,5 @@ def remove_error_handler( The error(s) for handlers to be removed. """ for handler in self.dispatcher.error_handlers: - if await handler.check_remove(error): + if handler.check_remove(error): self.dispatcher.error_handlers.remove(handler)