Exception handlers

Author: immmdreza

src/telegrambots/custom/__init__.py

@@ -4,6 +4,7 @@
 from .filters import callback_query as callback_query_filters
 from .contexts import (
     MessageContext,
+    TextMessageContext,
     CallbackQueryContext,
     ContinueWithInfo,
     ContextTemplate,
@@ -29,4 +30,5 @@
     "create_message_key",
     "create_key",
     "KeyBuilder",
+    "TextMessageContext"
 ]

src/telegrambots/custom/contexts/__init__.py

@@ -1,5 +1,5 @@
 from ._contexts.context_template import Context, ContextTemplate
-from ._contexts.message_context import MessageContext
+from ._contexts.message_context import MessageContext, TextMessageContext
 from ._contexts.callback_query_context import CallbackQueryContext
 from ._contexts._continuously_handler import ContinuouslyHandler, ContinueWithInfo
 
@@ -10,4 +10,5 @@
     "CallbackQueryContext",
     "ContinuouslyHandler",
     "ContinueWithInfo",
+    "TextMessageContext",
 ]

src/telegrambots/custom/contexts/_contexts/message_context.py

@@ -85,3 +85,18 @@ async def reply_text(
             allow_sending_without_reply=allow_sending_without_reply,
             reply_markup=reply_markup,
         )
+
+
+class TextMessageContext(MessageContext):
+    def __init__(
+        self, dp: "Dispatcher", update: Update, handler_tag: str, **kwargs: Any
+    ) -> None:
+        super().__init__(dp, update=update, handler_tag=handler_tag, **kwargs)
+
+    @final
+    @property
+    def text(self) -> str:
+        """Returns message text."""
+        if self.update.text is not None:
+            return self.update.text
+        raise ValueError("Update has no message")

src/telegrambots/custom/dispatcher.py

@@ -2,8 +2,6 @@
 import logging
 from typing import (
     Any,
-    Callable,
-    Coroutine,
     Mapping,
     Optional,
     cast,
@@ -20,6 +18,7 @@
 from .handlers._handlers.handler_template import HandlerTemplate
 from .processor import ProcessorTemplate, SequentialProcessor
 from .extensions.dispatcher import AddExtensions
+from .handlers import AbstractExceptionHandler, default_exception_handler
 
 if TYPE_CHECKING:
     from .client import TelegramBot
@@ -38,9 +37,6 @@ def __init__(
         self,
         _bot: "TelegramBot",
         *,
-        handle_error: Optional[
-            Callable[["TelegramBot", Exception], Coroutine[Any, Any, None]]
-        ] = None,
         processor_type: Optional[type[ProcessorTemplate[Update]]] = None,
     ) -> None:
 
@@ -54,7 +50,7 @@ def __init__(
         self._bot = _bot
         self._handlers: dict[type[Any], dict[str, HandlerTemplate]] = {}
         self._continuously_handlers: list[tuple[ContinuouslyHandlerTemplate]] = []
-        self._handle_error = handle_error
+        self._handle_errors: list[AbstractExceptionHandler] = []
 
         self._processor: ProcessorTemplate[Update]
         if processor_type is None:
@@ -124,6 +120,18 @@ def add_handler(self, tag: str, handler: HandlerTemplate):
         self._handlers[handler.update_type][tag] = handler
         dispatcher_logger.info(f"Added handler {handler.update_type.__name__}:{tag}")
 
+    def add_exception_handler(self, exception_handler: AbstractExceptionHandler):
+        """Adds an exception handler to the dispatcher.
+
+        Args:
+            exception_handler (`AbstractExceptionHandler`): The exception handler to add.
+        """
+        self._handle_errors.append(exception_handler)
+
+    def add_default_exception_handler(self):
+        """Adds the default exception handler to the dispatcher."""
+        self.add_exception_handler(default_exception_handler)
+
     @overload
     def add_continuously_handler(
         self,
@@ -251,5 +259,5 @@ async def _do_handling(
         return None
 
     async def _try_handle_error(self, e: Exception):
-        if self._handle_error is not None:
-            await self._handle_error(self._bot, e)
+        for handler in self._handle_errors:
+            await handler.try_handle(self, e)

src/telegrambots/custom/extensions/context.py

@@ -24,7 +24,7 @@
     CallbackQueryMessageId,
     CallbackQuerySenderId,
 )
-from ..filters import Filter
+from ..filters import Filter, messages as mf
 
 if TYPE_CHECKING:
     from .. import ContextTemplate
@@ -63,6 +63,7 @@ def callback_query(
         keys: Sequence[AbstractKeyResolver[CallbackQuery, Any]],
         filter: Filter[CallbackQuery],
         tag: Optional[str] = None,
+        other_continue_with: Optional[list[str]] = None,
         *args: Any,
         **kwargs: Any,
     ):
@@ -72,7 +73,10 @@ def decorator(
             _tag = tag or _function.__name__
             if not self._context.dp.handler_tag_exists(_tag, CallbackQuery):
                 self._context.dp.add.handlers.callback_query(
-                    _tag, _function, filter, [self._context.handler_tag]
+                    _tag,
+                    _function,
+                    filter,
+                    [self._context.handler_tag] + (other_continue_with or []),
                 )
             self._context.continue_with.callback_query(_tag, keys, *args, **kwargs)
 
@@ -83,6 +87,7 @@ def callback_query_form(
         user_id: int,
         filter: Filter[CallbackQuery],
         tag: Optional[str] = None,
+        other_continue_with: Optional[list[str]] = None,
         *args: Any,
         **kwargs: Any,
     ):
@@ -92,7 +97,10 @@ def decorator(
             _tag = tag or _function.__name__
             if not self._context.dp.handler_tag_exists(_tag, CallbackQuery):
                 self._context.dp.add.handlers.callback_query(
-                    _tag, _function, filter, [self._context.handler_tag]
+                    _tag,
+                    _function,
+                    filter,
+                    [self._context.handler_tag] + (other_continue_with or []),
                 )
             self._context.continue_with.callback_query(
                 _tag, [CallbackQuerySenderId(user_id)], *args, **kwargs
@@ -106,6 +114,7 @@ def callback_query_same_message_form(
         user_id: int,
         filter: Filter[CallbackQuery],
         tag: Optional[str] = None,
+        other_continue_with: Optional[list[str]] = None,
         *args: Any,
         **kwargs: Any,
     ):
@@ -115,7 +124,10 @@ def decorator(
             _tag = tag or _function.__name__
             if not self._context.dp.handler_tag_exists(_tag, CallbackQuery):
                 self._context.dp.add.handlers.callback_query(
-                    _tag, _function, filter, [self._context.handler_tag]
+                    _tag,
+                    _function,
+                    filter,
+                    [self._context.handler_tag] + (other_continue_with or []),
                 )
             self._context.continue_with.callback_query(
                 _tag,
@@ -131,6 +143,7 @@ def callback_query_same_message(
         message_id: int,
         filter: Filter[CallbackQuery],
         tag: Optional[str] = None,
+        other_continue_with: Optional[list[str]] = None,
         *args: Any,
         **kwargs: Any,
     ):
@@ -140,7 +153,10 @@ def decorator(
             _tag = tag or _function.__name__
             if not self._context.dp.handler_tag_exists(_tag, CallbackQuery):
                 self._context.dp.add.handlers.callback_query(
-                    _tag, _function, filter, [self._context.handler_tag]
+                    _tag,
+                    _function,
+                    filter,
+                    [self._context.handler_tag] + (other_continue_with or []),
                 )
             self._context.continue_with.callback_query(
                 _tag,
@@ -156,6 +172,7 @@ def message(
         keys: list[AbstractKeyResolver[Message, Any]],
         filter: "Filter[Message]",
         tag: Optional[str] = None,
+        other_continue_with: Optional[list[str]] = None,
         *args: Any,
         **kwargs: Any,
     ):
@@ -165,7 +182,10 @@ def decorator(
             _tag = tag or _function.__name__
             if not self._context.dp.handler_tag_exists(_tag, Message):
                 self._context.dp.add.handlers.message(
-                    _tag, _function, filter, [self._context.handler_tag]
+                    _tag,
+                    _function,
+                    filter,
+                    [self._context.handler_tag] + (other_continue_with or []),
                 )
             self._context.continue_with.message(_tag, keys, *args, **kwargs)
 
@@ -176,6 +196,7 @@ def message_from(
         user_id: int,
         filter: "Filter[Message]",
         tag: Optional[str] = None,
+        other_continue_with: Optional[list[str]] = None,
         *args: Any,
         **kwargs: Any,
     ):
@@ -185,7 +206,35 @@ def decorator(
             _tag = tag or _function.__name__
             if not self._context.dp.handler_tag_exists(_tag, Message):
                 self._context.dp.add.handlers.message(
-                    _tag, _function, filter, [self._context.handler_tag]
+                    _tag,
+                    _function,
+                    filter,
+                    [self._context.handler_tag] + (other_continue_with or []),
+                )
+            self._context.continue_with.message(
+                _tag, [MessageSenderId(user_id)], *args, **kwargs
+            )
+
+        return decorator
+
+    def text_input_from(
+        self,
+        user_id: int,
+        tag: Optional[str] = None,
+        other_continue_with: Optional[list[str]] = None,
+        *args: Any,
+        **kwargs: Any,
+    ):
+        def decorator(
+            _function: Callable[["MessageContext"], Coroutine[Any, Any, None]]
+        ):
+            _tag = tag or _function.__name__
+            if not self._context.dp.handler_tag_exists(_tag, Message):
+                self._context.dp.add.handlers.message(
+                    _tag,
+                    _function,  # type: ignore
+                    mf.text_message,
+                    [self._context.handler_tag] + (other_continue_with or []),
                 )
             self._context.continue_with.message(
                 _tag, [MessageSenderId(user_id)], *args, **kwargs

src/telegrambots/custom/handlers/__init__.py

@@ -1,5 +1,17 @@
 from ._handlers.handler_template import Handler
 from ._handlers.update_handler import CallbackQueryHandler, MessageHandler
+from .exceptions.exception_handler import (
+    AbstractExceptionHandler,
+    ExceptionHandler,
+    default_exception_handler,
+)
 
 
-__all__ = ["Handler", "CallbackQueryHandler", "MessageHandler"]
+__all__ = [
+    "Handler",
+    "CallbackQueryHandler",
+    "MessageHandler",
+    "AbstractExceptionHandler",
+    "ExceptionHandler",
+    "default_exception_handler",
+]

src/telegrambots/custom/handlers/exceptions/__init__.py

@@ -0,0 +1,47 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Callable, Coroutine
+
+if TYPE_CHECKING:
+    from ...dispatcher import Dispatcher
+
+
+class AbstractExceptionHandler(ABC):
+    def __init__(self, type: type[Exception], exact_exception: bool = False) -> None:
+        super().__init__()
+        self._exact = exact_exception
+        self._type = type
+
+    @abstractmethod
+    async def _handle_exception(self, dp: "Dispatcher", e: Exception):
+        ...
+
+    def should_handle(self, e: Exception):
+        if self._exact:
+            return isinstance(e, self._type)
+        else:
+            return issubclass(type(e), self._type)
+
+    async def try_handle(self, dp: "Dispatcher", e: Exception):
+        if self.should_handle(e):
+            await self._handle_exception(dp, e)
+
+
+class ExceptionHandler(AbstractExceptionHandler):
+    def __init__(
+        self,
+        handle: Callable[["Dispatcher", Exception], Coroutine[Any, Any, None]],
+        type: type[Exception],
+        exact_exception: bool = False,
+    ) -> None:
+        super().__init__(type, exact_exception)
+        self._handle = handle
+
+    async def _handle_exception(self, dp: "Dispatcher", e: Exception):
+        return await self._handle(dp, e)
+
+
+async def __handle_exception(dp: "Dispatcher", e: Exception):
+    print(e)
+
+
+default_exception_handler = ExceptionHandler(__handle_exception, Exception)