Add Error Handler support for handling unexpected update handler exceptions (#285)

Author: zrekryu

pyrogram/dispatcher.py

@@ -17,14 +17,16 @@
 #  along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
 import asyncio
+from collections import OrderedDict
 import inspect
 import logging
-from collections import OrderedDict
+from typing import Dict
 
 import pyrogram
 from pyrogram import utils
 from pyrogram.handlers import (
-    CallbackQueryHandler, MessageHandler, EditedMessageHandler, DeletedMessagesHandler,
+    Handler,
+    ErrorHandler, CallbackQueryHandler, MessageHandler, EditedMessageHandler, DeletedMessagesHandler,
     UserStatusHandler, RawUpdateHandler, InlineQueryHandler, PollHandler, PreCheckoutQueryHandler,
     ChosenInlineResultHandler, ChatMemberUpdatedHandler, ChatJoinRequestHandler, StoryHandler,
     ShippingQueryHandler, MessageReactionHandler, MessageReactionCountHandler, ChatBoostHandler,
@@ -294,7 +296,7 @@ async def stop(self, clear_handlers: bool = True):
 
             log.info("Stopped %s HandlerTasks", self.client.workers)
 
-    def add_handler(self, handler, group: int):
+    def add_handler(self, handler: Handler, group: int):
         async def fn():
             for lock in self.locks_list:
                 await lock.acquire()
@@ -311,16 +313,21 @@ async def fn():
 
         self.client.loop.create_task(fn())
 
-    def remove_handler(self, handler, group: int):
+    def remove_handler(self, handler: Handler, group: int):
         async def fn():
             for lock in self.locks_list:
                 await lock.acquire()
 
             try:
                 if group not in self.groups:
-                    raise ValueError(f"Group {group} does not exist. Handler was not removed.")
+                    raise ValueError(
+                        f"Group {group} does not exist. Handler was not removed."
+                    )
 
                 self.groups[group].remove(handler)
+
+                if not self.groups[group]:
+                    del self.groups[group]
             finally:
                 for lock in self.locks_list:
                     lock.release()
@@ -347,6 +354,9 @@ async def handler_worker(self, lock):
                 async with lock:
                     for group in self.groups.values():
                         for handler in group:
+                            if isinstance(handler, ErrorHandler):
+                                continue
+
                             args = None
 
                             if isinstance(handler, handler_type):
@@ -382,11 +392,62 @@ async def handler_worker(self, lock):
                                 raise
                             except pyrogram.ContinuePropagation:
                                 continue
-                            except Exception as e:
-                                log.exception(e)
+                            except Exception as exc:
+                                await self.handle_update_handler_exception(
+                                    exc, handler, update, users, chats
+                                )
 
                             break
             except pyrogram.StopPropagation:
                 pass
             except Exception as e:
                 log.exception(e)
+
+    async def handle_update_handler_exception(
+        self,
+        exc: Exception,
+        update_handler: Handler,
+        update: "pyrogram.raw.base.Update",
+        users: Dict[int, "pyrogram.raw.base.User"],
+        chats: Dict[int, "pyrogram.raw.base.Chat"]
+    ) -> None:
+        handled = False
+        try:
+            for group in self.groups.values():
+                for handler in group:
+                    if not isinstance(handler, ErrorHandler):
+                        continue
+
+                    if not isinstance(exc, handler.exceptions):
+                        continue
+
+                    try:
+                        if inspect.iscoroutinefunction(handler.callback):
+                            await handler.callback(
+                                self.client, exc, update_handler, update, users, chats
+                            )
+                        else:
+                            await self.client.loop.run_in_executor(
+                                self.client.executor, handler.callback,
+                                self.client, exc, update_handler, update, users, chats
+                            )
+                    except pyrogram.StopPropagation:
+                        handled = True
+                        raise
+                    except pyrogram.ContinuePropagation:
+                        handled = True
+                        continue
+                    except Exception:
+                        log.exception("Error handler raised an exception:")
+                    else:
+                        handled = True
+
+                    break
+        except pyrogram.StopPropagation:
+            pass
+        finally:
+            if not handled:
+                log.error(
+                    f"Unexpected exception raised in {type(update_handler).__name__}:",
+                    exc_info=(type(exc), exc, exc.__traceback__)
+                )

pyrogram/handlers/__init__.py

@@ -17,6 +17,7 @@
 #  along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
 from .handler import Handler
+from .error_handler import ErrorHandler
 from .business_connection_handler import BusinessConnectionHandler
 from .business_message_handler import BusinessMessageHandler
 from .callback_query_handler import CallbackQueryHandler

pyrogram/handlers/error_handler.py

@@ -0,0 +1,89 @@
+#  Pyrogram - Telegram MTProto API Client Library for Python
+#  Copyright (C) 2017-present Dan <https://github.com/delivrance>
+#
+#  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 <http://www.gnu.org/licenses/>.
+
+from typing import Callable, Optional, Sequence, Union
+
+import pyrogram
+from pyrogram.filters import Filter
+from pyrogram.types import Update
+
+from .handler import Handler
+
+
+class ErrorHandler(Handler):
+    """The Error handler class. Used to handle unexpected errors.
+
+    It is intended to be used with :meth:`~pyrogram.Client.add_handler`.
+
+    For a more convenient way to register this handler, see the
+    :meth:`~pyrogram.Client.on_error` decorator.
+
+    Parameters:
+        callback (``Callable``):
+            A function that will be called whenever an unexpected error is raised.
+            It takes the following positional arguments: *(exception, handler, client, *args)*.
+
+        exceptions (``Exception`` | List of ``Exception``, *optional*):
+            An exception type or a sequence of exception types that this handler should handle.
+            If None, the handler will catch any exception that is a subclass of ``Exception``.
+
+        filters (:obj:`Filter`, *optional*):
+            Pass one or more filters to allow only a subset of updates to be passed
+            in your callback function.
+
+    Other parameters passed to the callback:
+        client (:obj:`~pyrogram.Client`):
+            The Client instance, useful when calling other API methods inside the error handler.
+
+        exception (``Exception``):
+            The Exception instance that was raised.
+
+        handler (:obj:`~pyrogram.handlers.handler.Handler`):
+            The Handler instance from which the exception was raised.
+
+        update (:obj:`~pyrogram.raw.base.Update`):
+            The received update, which can be one of the many single Updates listed in the
+            :obj:`~pyrogram.raw.base.Update` base type.
+
+        users (``dict``):
+            Dictionary of all :obj:`~pyrogram.types.User` mentioned in the update.
+            You can access extra info about the user (such as *first_name*, *last_name*, etc...) by using
+            the IDs you find in the *update* argument (e.g.: *users[1768841572]*).
+
+        chats (``dict``):
+            Dictionary of all :obj:`~pyrogram.types.Chat` and
+            :obj:`~pyrogram.raw.types.Channel` mentioned in the update.
+            You can access extra info about the chat (such as *title*, *participants_count*, etc...)
+            by using the IDs you find in the *update* argument (e.g.: *chats[1701277281]*).
+
+    """
+
+    def __init__(
+        self,
+        callback: Callable,
+        exceptions: Optional[Union[Exception, Sequence[Exception]]] = None,
+        filters: Optional[Filter] = None,
+    ):
+        super().__init__(callback, filters)
+
+        if exceptions is None:
+            self.exceptions = (Exception,)
+        elif isinstance(exceptions, Sequence):
+            self.exceptions = tuple(exceptions)
+        else:
+            self.exceptions = (exceptions,)

pyrogram/methods/decorators/__init__.py

@@ -16,6 +16,7 @@
 #  You should have received a copy of the GNU Lesser General Public License
 #  along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
+from .on_error import OnError
 from .on_business_connection import OnBusinessConnection
 from .on_business_message import OnBusinessMessage
 from .on_callback_query import OnCallbackQuery
@@ -45,6 +46,7 @@
 
 
 class Decorators(
+    OnError,
     OnBusinessConnection,
     OnBusinessMessage,
     OnMessage,

pyrogram/methods/decorators/on_error.py

@@ -0,0 +1,65 @@
+#  Pyrogram - Telegram MTProto API Client Library for Python
+#  Copyright (C) 2017-present Dan <https://github.com/delivrance>
+#
+#  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 <http://www.gnu.org/licenses/>.
+
+from typing import Callable, Optional, Sequence, Union
+
+import pyrogram
+from pyrogram.filters import Filter
+
+
+class OnError:
+    def on_error(
+        self: Optional[Union["OnError", Filter]] = None,
+        exceptions: Optional[Union[Exception, Sequence[Exception]]] = None,
+        filters: Optional[Filter] = None,
+        group: int = 0,
+    ) -> Callable:
+        """Decorator for handling unexpected errors.
+
+        This does the same thing as :meth:`~pyrogram.Client.add_handler` using the
+        :obj:`~pyrogram.handlers.ErrorHandler`.
+
+        .. include:: /_includes/usable-by/users-bots.rst
+
+        Parameters:
+            exceptions (``Exception`` | List of ``Exception``, *optional*):
+                An exception type or a sequence of exception types that this handler should handle.
+                If None, the handler will catch any exception that is a subclass of ``Exception``.
+
+            filters (:obj:`~pyrogram.filters`, *optional*):
+                Pass one or more filters to allow only a subset of messages to be passed
+                in your function.
+
+            group (``int``, *optional*):
+                The group identifier, defaults to 0.
+        """
+
+        def decorator(func: Callable) -> Callable:
+            if isinstance(self, pyrogram.Client):
+                self.add_handler(pyrogram.handlers.ErrorHandler(func, filters, exceptions), group)
+            else:
+                if not hasattr(func, "handlers"):
+                    func.handlers = []
+
+                func.handlers.append(
+                    (pyrogram.handlers.ErrorHandler(func, filters, exceptions), group)
+                )
+
+            return func
+
+        return decorator