pyrofork: Add Error handler

Signed-off-by: Yasir Aris M <git@yasirdev.my.id>

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()

pyrogram/dispatcher.py

@@ -23,16 +23,16 @@
 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,
   CallbackQueryHandler,
   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
+            )

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

pyrogram/handlers/error_handler.py

@@ -0,0 +1,79 @@
+#  Pyrofork - Telegram MTProto API Client Library for Python
+#  Copyright (C) 2017-present Dan <https://github.com/delivrance>
+#  Copyright (C) 2022-present Mayuri-Chan <https://github.com/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 <http://www.gnu.org/licenses/>.
+
+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)

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,

pyrogram/methods/decorators/on_error.py

@@ -0,0 +1,50 @@
+#  Pyrofork - Telegram MTProto API Client Library for Python
+#  Copyright (C) 2017-present Dan <https://github.com/delivrance>
+#  Copyright (C) 2022-present Mayuri-Chan <https://github.com/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 <http://www.gnu.org/licenses/>.
+
+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

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,