python-telegram-bot · tsnoam · Aug 12, 2017 · Aug 12, 2017 · Aug 12, 2017 · Eldinnie
diff --git a/telegram/ext/__init__.py b/telegram/ext/__init__.py
@@ -18,7 +18,7 @@
 # along with this program.  If not, see [http://www.gnu.org/licenses/].
 """Extensions over the Telegram Bot API to facilitate bot making"""
 
-from .dispatcher import Dispatcher, DispatcherHandlerContinue, DispatcherHandlerStop, run_async
+from .dispatcher import Dispatcher, DispatcherHandlerStop, run_async
 from .jobqueue import JobQueue, Job
 from .updater import Updater
 from .callbackqueryhandler import CallbackQueryHandler
@@ -43,4 +43,4 @@
            'MessageHandler', 'BaseFilter', 'Filters', 'RegexHandler', 'StringCommandHandler',
            'StringRegexHandler', 'TypeHandler', 'ConversationHandler',
            'PreCheckoutQueryHandler', 'ShippingQueryHandler', 'MessageQueue', 'DelayQueue',
-           'DispatcherHandlerContinue', 'DispatcherHandlerStop', 'run_async')
+           'DispatcherHandlerStop', 'run_async')
diff --git a/telegram/ext/dispatcher.py b/telegram/ext/dispatcher.py
@@ -55,25 +55,8 @@ def async_func(*args, **kwargs):
     return async_func
 
 
-class DispatcherHandlerFlow(Exception):
-    """
-    Dispatcher update processing manipulation exceptions are base on this class.
-    """
-    pass
-
-
-class DispatcherHandlerContinue(DispatcherHandlerFlow):
-    """
-    If check Handler's check_updated returned true, but execution of handler raised this,
-    then handlers checking will continue.
-    """
-    pass
-
-
-class DispatcherHandlerStop(DispatcherHandlerFlow):
-    """
-    Raise this in handler to prevent execution any other handlers (even in different group).
-    """
+class DispatcherHandlerStop(Exception):
+    """Raise this in handler to prevent execution any other handler (even in different group)."""
     pass
 
 
@@ -178,9 +161,10 @@ def _pooled(self):
                 break
 
             promise.run()
-            if isinstance(promise.exception, DispatcherHandlerFlow):
-                self.logger.warning('DispatcherHandlerFlow is not supported with async '
-                                    'functions; func: %s', promise.pooled_function.__name__)
+            if isinstance(promise.exception, DispatcherHandlerStop):
+                self.logger.warning(
+                    'DispatcherHandlerStop is not supported with async functions; func: %s',
+                    promise.pooled_function.__name__)
 
     def run_async(self, func, *args, **kwargs):
         """
@@ -284,63 +268,54 @@ def process_update(self, update):
             return
         for group in self.groups:
             try:
-                for handler in self.handlers[group]:
-                    try:
-                        if handler.check_update(update):
-                            try:
-                                handler.handle_update(update, self)
-                            except DispatcherHandlerContinue:
-                                continue
-                            break
-                    except DispatcherHandlerFlow:
-                        raise
-                    except TelegramError as te:
-                        self.logger.warning('A TelegramError was raised while processing the '
-                                            'Update.')
-
-                        try:
-                            self.dispatch_error(update, te)
-                        except Exception:
-                            self.logger.exception('An uncaught error was raised while '
-                                                  'handling the error')
-                        finally:
-                            break
-
-                    # Errors should not stop the thread
-                    except Exception:
-                        self.logger.exception('An uncaught error was raised while '
-                                              'processing the update')
-                        break
+                for handler in (x for x in self.handlers[group] if x.check_update(update)):
+                    handler.handle_update(update, self)
+                    break
+
+            # Stop processing with any other handler.
             except DispatcherHandlerStop:
+                self.logger.debug('Stopping further handlers due to DispatcherHandlerStop')
                 break
 
+            # Dispatch any error.
+            except TelegramError as te:
+                self.logger.warning('A TelegramError was raised while processing the Update')
+
+                try:
+                    self.dispatch_error(update, te)
+                except DispatcherHandlerStop:
+                    self.logger.debug('Error handler stopped further handlers')
+                    break
+                except Exception:
+                    self.logger.exception('An uncaught error was raised while handling the error')
+
+            # Errors should not stop the thread.
+            except Exception:
+                self.logger.exception('An uncaught error was raised while processing the update')
+
     def add_handler(self, handler, group=DEFAULT_GROUP):
-        """
-        Register a handler.
+        """Register a handler.
 
-        TL;DR: Order and priority counts. 0 or 1 handlers per group will be
-        used.
+        TL;DR: Order and priority counts. 0 or 1 handlers per group will be used.
 
         A handler must be an instance of a subclass of :class:`telegram.ext.Handler`. All handlers
         are organized in groups with a numeric value. The default group is 0. All groups will be
-        evaluated for handling an update, but only 0 or 1 handler per group will be used,
-        except situations when :class:`telegram.DispatcherHandlerContinue` or
-        :class:`telegram.DispatcherHandlerStop` were raised.
+        evaluated for handling an update, but only 0 or 1 handler per group will be used. If
+        :class:`telegram.DispatcherHandlerStop` is raised from one of the handlers, no further
+        handlers (regardless of the group) will be called.
 
         The priority/order of handlers is determined as follows:
 
           * Priority of the group (lower group number == higher priority)
-          * The first handler in a group which should handle an update will be
-            used. Other handlers from the group will not be used. The order in
-            which handlers were added to the group defines the priority.
-          * If :class:`telegram.DispatcherHandlerContinue` was raised, then next handler in the
-            same group will be called.
-          * If :class:`telegram.DispatcherHandlerStop` was raised, then zero handlers (even
-            from other groups) will called.
+          * The first handler in a group which should handle an update (see
+            :method:`telegram.ext.Handler.check_update`) will be used. Other handlers from the
+            group will not be used. The order in which handlers were added to the group defines the
+            priority.
 
         Args:
             handler (:class:`telegram.ext.Handler`): A Handler instance.
             group (:obj:`int`, optional): The group identifier. Default is 0.
+
         """
 
         if not isinstance(handler, Handler):

diff --git a/tests/test_dispatcher.py b/tests/test_dispatcher.py
@@ -24,8 +24,7 @@
 
 from telegram import TelegramError, Message, User, Chat, Update
 from telegram.ext import MessageHandler, Filters, CommandHandler
-from telegram.ext.dispatcher import run_async, Dispatcher, DispatcherHandlerContinue, \
-    DispatcherHandlerStop
+from telegram.ext.dispatcher import run_async, Dispatcher, DispatcherHandlerStop
 from tests.conftest import create_dp
 
 
@@ -173,12 +172,12 @@ def test_add_handler_errors(self, dp):
         with pytest.raises(TypeError, match='group is not int'):
             dp.add_handler(handler, 'one')
 
-    def test_handler_flow_continue(self, bot, dp):
+    def test_flow_stop(self, dp, bot):
         passed = []
 
         def start1(b, u):
             passed.append('start1')
-            raise DispatcherHandlerContinue
+            raise DispatcherHandlerStop
 
         def start2(b, u):
             passed.append('start2')
@@ -192,19 +191,20 @@ def error(b, u, e):
 
         update = Update(1, message=Message(1, None, None, None, text='/start', bot=bot))
 
-        # If Continue raised next handler should be proceed.
+        # If Stop raised handlers in other groups should not be called.
         passed = []
-        dp.add_handler(CommandHandler('start', start1))
-        dp.add_handler(CommandHandler('start', start2))
+        dp.add_handler(CommandHandler('start', start1), 1)
+        dp.add_handler(CommandHandler('start', start3), 1)
+        dp.add_handler(CommandHandler('start', start2), 2)
         dp.process_update(update)
-        assert passed == ['start1', 'start2']
+        assert passed == ['start1']
 
-    def test_dispatcher_handler_flow_stop(self, dp, bot):
+    def test_exception_in_handler(self, dp, bot):
         passed = []
 
         def start1(b, u):
             passed.append('start1')
-            raise DispatcherHandlerStop
+            raise Exception('General exception')
 
         def start2(b, u):
             passed.append('start2')
@@ -218,10 +218,73 @@ def error(b, u, e):
 
         update = Update(1, message=Message(1, None, None, None, text='/start', bot=bot))
 
-        # If Stop raised handlers in other groups should not be called.
+        # If an unhandled exception was caught, no further handlers from the same group should be
+        # called.
         passed = []
         dp.add_handler(CommandHandler('start', start1), 1)
-        dp.add_handler(CommandHandler('start', start3), 1)
-        dp.add_handler(CommandHandler('start', start2), 2)
+        dp.add_handler(CommandHandler('start', start2), 1)
+        dp.add_handler(CommandHandler('start', start3), 2)
+        dp.add_error_handler(error)
         dp.process_update(update)
-        assert passed == ['start1']
+        assert passed == ['start1', 'start3']
+
+    def test_telegram_error_in_handler(self, dp, bot):
+        passed = []
+        err = TelegramError('Telegram error')
+
+        def start1(b, u):
+            passed.append('start1')
+            raise err
+
+        def start2(b, u):
+            passed.append('start2')
+
+        def start3(b, u):
+            passed.append('start3')
+
+        def error(b, u, e):
+            passed.append('error')
+            passed.append(e)
+
+        update = Update(1, message=Message(1, None, None, None, text='/start', bot=bot))
+
+        # If a TelegramException was caught, an error handler should be called and no further
+        # handlers from the same group should be called.
+        dp.add_handler(CommandHandler('start', start1), 1)
+        dp.add_handler(CommandHandler('start', start2), 1)
+        dp.add_handler(CommandHandler('start', start3), 2)
+        dp.add_error_handler(error)
+        dp.process_update(update)
+        assert passed == ['start1', 'error', err, 'start3']
+        assert passed[2] is err
+
+    def test_flow_stop_in_error_handler(self, dp, bot):
+        passed = []
+        err = TelegramError('Telegram error')
+
+        def start1(b, u):
+            passed.append('start1')
+            raise err
+
+        def start2(b, u):
+            passed.append('start2')
+
+        def start3(b, u):
+            passed.append('start3')
+
+        def error(b, u, e):
+            passed.append('error')
+            passed.append(e)
+            raise DispatcherHandlerStop
+
+        update = Update(1, message=Message(1, None, None, None, text='/start', bot=bot))
+
+        # If a TelegramException was caught, an error handler should be called and no further
+        # handlers from the same group should be called.
+        dp.add_handler(CommandHandler('start', start1), 1)
+        dp.add_handler(CommandHandler('start', start2), 1)
+        dp.add_handler(CommandHandler('start', start3), 2)
+        dp.add_error_handler(error)
+        dp.process_update(update)
+        assert passed == ['start1', 'error', err]
+        assert passed[2] is err