From 6aacde189ec6096be42af8d5bda006f2c64ba4ac Mon Sep 17 00:00:00 2001 From: Ihor Polyakov Date: Sun, 30 Jul 2017 00:15:43 +0700 Subject: [PATCH] Flow control ability in Dispatcher (#738) fixes #666 --- telegram/ext/__init__.py | 5 ++- telegram/ext/dispatcher.py | 62 +++++++++++++++++++++++++------ telegram/utils/promise.py | 4 ++ tests/test_updater.py | 75 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 132 insertions(+), 14 deletions(-) diff --git a/telegram/ext/__init__.py b/telegram/ext/__init__.py index 472b953f2..fd0ca2870 100644 --- 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 +from .dispatcher import Dispatcher, DispatcherHandlerContinue, DispatcherHandlerStop, run_async from .jobqueue import JobQueue, Job from .updater import Updater from .callbackqueryhandler import CallbackQueryHandler @@ -42,4 +42,5 @@ __all__ = ('Dispatcher', 'JobQueue', 'Job', 'Updater', 'CallbackQueryHandler', 'ChosenInlineResultHandler', 'CommandHandler', 'Handler', 'InlineQueryHandler', 'MessageHandler', 'BaseFilter', 'Filters', 'RegexHandler', 'StringCommandHandler', 'StringRegexHandler', 'TypeHandler', 'ConversationHandler', - 'PreCheckoutQueryHandler', 'ShippingQueryHandler', 'MessageQueue', 'DelayQueue') + 'PreCheckoutQueryHandler', 'ShippingQueryHandler', 'MessageQueue', 'DelayQueue', + 'DispatcherHandlerContinue', 'DispatcherHandlerStop', 'run_async') diff --git a/telegram/ext/dispatcher.py b/telegram/ext/dispatcher.py index 2b9a27c1e..8d79dd27b 100644 --- a/telegram/ext/dispatcher.py +++ b/telegram/ext/dispatcher.py @@ -55,6 +55,28 @@ def run_async(func): 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). + """ + pass + + class Dispatcher(object): """ This class dispatches all kinds of updates to its registered handlers. @@ -162,6 +184,9 @@ class Dispatcher(object): break promise.run() + if isinstance(promise.exception, DispatcherHandlerFlow): + self.logger.warning('DispatcherHandlerFlow is not supported with async ' + 'functions; func: %s', promise.pooled_function.__name__) def run_async(self, func, *args, **kwargs): """ @@ -255,24 +280,29 @@ class Dispatcher(object): Processes a single update. Args: - update (:obj:`str` | :class:`telegram.Update`): The update to process. + update (:obj:`str` | :class:`telegram.Update` | :class:`telegram.TelegramError`): + The update to process. """ # An error happened while polling if isinstance(update, TelegramError): self.dispatch_error(None, update) - - else: - for group in self.groups: + return + for group in self.groups: + try: for handler in self.handlers[group]: try: if handler.check_update(update): - handler.handle_update(update, self) + try: + handler.handle_update(update, self) + except DispatcherHandlerContinue: + continue break - # Dispatch any errors + except DispatcherHandlerFlow: + raise except TelegramError as te: - self.logger.warn('A TelegramError was raised while processing the ' - 'Update.') + self.logger.warning('A TelegramError was raised while processing the ' + 'Update.') try: self.dispatch_error(update, te) @@ -287,6 +317,8 @@ class Dispatcher(object): self.logger.exception('An uncaught error was raised while ' 'processing the update') break + except DispatcherHandlerStop: + break def add_handler(self, handler, group=DEFAULT_GROUP): """ @@ -297,7 +329,9 @@ class Dispatcher(object): 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. + 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. The priority/order of handlers is determined as follows: @@ -305,6 +339,10 @@ class Dispatcher(object): * 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. Args: handler (:class:`telegram.ext.Handler`): A Handler instance. @@ -343,7 +381,7 @@ class Dispatcher(object): Registers an error handler in the Dispatcher. Args: - handler (:obj:`callable`): A function that takes ``Bot, Update, TelegramError`` as + callback (:obj:`callable`): A function that takes ``Bot, Update, TelegramError`` as arguments. """ @@ -354,7 +392,7 @@ class Dispatcher(object): Removes an error handler. Args: - handler (:obj:`callable`): The error handler to remove. + callback (:obj:`callable`): The error handler to remove. """ if callback in self.error_handlers: @@ -365,7 +403,7 @@ class Dispatcher(object): Dispatches an error. Args: - update (:obj:`str` | :class:`telegram.Update`): The update that caused the error + update (:obj:`str` | :class:`telegram.Update` | None): The update that caused the error error (:class:`telegram.TelegramError`): The Telegram error that was raised. """ diff --git a/telegram/utils/promise.py b/telegram/utils/promise.py index 43e2e7599..d3d5343a6 100644 --- a/telegram/utils/promise.py +++ b/telegram/utils/promise.py @@ -56,3 +56,7 @@ class Promise(object): if self._exception is not None: raise self._exception # pylint: disable=raising-bad-type return self._result + + @property + def exception(self): + return self._exception diff --git a/tests/test_updater.py b/tests/test_updater.py index 4f23322f5..07849cd27 100644 --- a/tests/test_updater.py +++ b/tests/test_updater.py @@ -939,6 +939,81 @@ class UpdaterTest(BaseTest, unittest.TestCase): def test_noTokenOrBot(self): self.assertRaises(ValueError, Updater) + def test_dispatcher_handler_flow_continue(self): + passed = [] + + def start1(b, u): + passed.append('start1') + raise DispatcherHandlerContinue + + 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) + + # noinspection PyTypeChecker + update = Update(1, message=Message(1, None, None, None, text='/start', bot=self._bot)) + + # Without raising Continue everything should work as before + passed = [] + dp = Dispatcher(self._bot, Queue()) + dp.add_handler(CommandHandler('start', start3)) + dp.add_handler(CommandHandler('start', start2)) + dp.add_error_handler(error) + dp.process_update(update) + self.assertEqual(passed, ['start3']) + + # If Continue raised next handler should be proceed. + passed = [] + dp = Dispatcher(self._bot, Queue()) + dp.add_handler(CommandHandler('start', start1)) + dp.add_handler(CommandHandler('start', start2)) + dp.process_update(update) + self.assertEqual(passed, ['start1', 'start2']) + + def test_dispatcher_handler_flow_stop(self): + passed = [] + + def start1(b, u): + passed.append('start1') + raise DispatcherHandlerStop + + 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) + + # noinspection PyTypeChecker + update = Update(1, message=Message(1, None, None, None, text='/start', bot=self._bot)) + + # Without raising Stop everything should work as before + passed = [] + dp = Dispatcher(self._bot, Queue()) + dp.add_handler(CommandHandler('start', start3), 1) + dp.add_handler(CommandHandler('start', start2), 2) + dp.add_error_handler(error) + dp.process_update(update) + self.assertEqual(passed, ['start3', 'start2']) + + # If Stop raised handlers in other groups should not be called. + passed = [] + dp = Dispatcher(self._bot, Queue()) + 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) + self.assertEqual(passed, ['start1']) + class MockBot(object): def __init__(self,