Skip to content

Commit 6aacde1

Browse files
ihorutsnoam
ihoru
authored and
tsnoam
committed
Flow control ability in Dispatcher (python-telegram-bot#738)
fixes python-telegram-bot#666
1 parent 5d3f557 commit 6aacde1

File tree

4 files changed

+132
-14
lines changed

4 files changed

+132
-14
lines changed

‎telegram/ext/__init__.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
# along with this program. If not, see [http://www.gnu.org/licenses/].
1919
"""Extensions over the Telegram Bot API to facilitate bot making"""
2020

21-
from .dispatcher import Dispatcher
21+
from .dispatcher import Dispatcher, DispatcherHandlerContinue, DispatcherHandlerStop, run_async
2222
from .jobqueue import JobQueue, Job
2323
from .updater import Updater
2424
from .callbackqueryhandler import CallbackQueryHandler
@@ -42,4 +42,5 @@
4242
'ChosenInlineResultHandler', 'CommandHandler', 'Handler', 'InlineQueryHandler',
4343
'MessageHandler', 'BaseFilter', 'Filters', 'RegexHandler', 'StringCommandHandler',
4444
'StringRegexHandler', 'TypeHandler', 'ConversationHandler',
45-
'PreCheckoutQueryHandler', 'ShippingQueryHandler', 'MessageQueue', 'DelayQueue')
45+
'PreCheckoutQueryHandler', 'ShippingQueryHandler', 'MessageQueue', 'DelayQueue',
46+
'DispatcherHandlerContinue', 'DispatcherHandlerStop', 'run_async')

‎telegram/ext/dispatcher.py‎

Lines changed: 50 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,28 @@ def async_func(*args, **kwargs):
5555
return async_func
5656

5757

58+
class DispatcherHandlerFlow(Exception):
59+
"""
60+
Dispatcher update processing manipulation exceptions are base on this class.
61+
"""
62+
pass
63+
64+
65+
class DispatcherHandlerContinue(DispatcherHandlerFlow):
66+
"""
67+
If check Handler's check_updated returned true, but execution of handler raised this,
68+
then handlers checking will continue.
69+
"""
70+
pass
71+
72+
73+
class DispatcherHandlerStop(DispatcherHandlerFlow):
74+
"""
75+
Raise this in handler to prevent execution any other handlers (even in different group).
76+
"""
77+
pass
78+
79+
5880
class Dispatcher(object):
5981
"""
6082
This class dispatches all kinds of updates to its registered handlers.
@@ -162,6 +184,9 @@ def _pooled(self):
162184
break
163185

164186
promise.run()
187+
if isinstance(promise.exception, DispatcherHandlerFlow):
188+
self.logger.warning('DispatcherHandlerFlow is not supported with async '
189+
'functions; func: %s', promise.pooled_function.__name__)
165190

166191
def run_async(self, func, *args, **kwargs):
167192
"""
@@ -255,24 +280,29 @@ def process_update(self, update):
255280
Processes a single update.
256281
257282
Args:
258-
update (:obj:`str` | :class:`telegram.Update`): The update to process.
283+
update (:obj:`str` | :class:`telegram.Update` | :class:`telegram.TelegramError`):
284+
The update to process.
259285
"""
260286

261287
# An error happened while polling
262288
if isinstance(update, TelegramError):
263289
self.dispatch_error(None, update)
264-
265-
else:
266-
for group in self.groups:
290+
return
291+
for group in self.groups:
292+
try:
267293
for handler in self.handlers[group]:
268294
try:
269295
if handler.check_update(update):
270-
handler.handle_update(update, self)
296+
try:
297+
handler.handle_update(update, self)
298+
except DispatcherHandlerContinue:
299+
continue
271300
break
272-
# Dispatch any errors
301+
except DispatcherHandlerFlow:
302+
raise
273303
except TelegramError as te:
274-
self.logger.warn('A TelegramError was raised while processing the '
275-
'Update.')
304+
self.logger.warning('A TelegramError was raised while processing the '
305+
'Update.')
276306

277307
try:
278308
self.dispatch_error(update, te)
@@ -287,6 +317,8 @@ def process_update(self, update):
287317
self.logger.exception('An uncaught error was raised while '
288318
'processing the update')
289319
break
320+
except DispatcherHandlerStop:
321+
break
290322

291323
def add_handler(self, handler, group=DEFAULT_GROUP):
292324
"""
@@ -297,14 +329,20 @@ def add_handler(self, handler, group=DEFAULT_GROUP):
297329
298330
A handler must be an instance of a subclass of :class:`telegram.ext.Handler`. All handlers
299331
are organized in groups with a numeric value. The default group is 0. All groups will be
300-
evaluated for handling an update, but only 0 or 1 handler per group will be used.
332+
evaluated for handling an update, but only 0 or 1 handler per group will be used,
333+
except situations when :class:`telegram.DispatcherHandlerContinue` or
334+
:class:`telegram.DispatcherHandlerStop` were raised.
301335
302336
The priority/order of handlers is determined as follows:
303337
304338
* Priority of the group (lower group number == higher priority)
305339
* The first handler in a group which should handle an update will be
306340
used. Other handlers from the group will not be used. The order in
307341
which handlers were added to the group defines the priority.
342+
* If :class:`telegram.DispatcherHandlerContinue` was raised, then next handler in the
343+
same group will be called.
344+
* If :class:`telegram.DispatcherHandlerStop` was raised, then zero handlers (even
345+
from other groups) will called.
308346
309347
Args:
310348
handler (:class:`telegram.ext.Handler`): A Handler instance.
@@ -343,7 +381,7 @@ def add_error_handler(self, callback):
343381
Registers an error handler in the Dispatcher.
344382
345383
Args:
346-
handler (:obj:`callable`): A function that takes ``Bot, Update, TelegramError`` as
384+
callback (:obj:`callable`): A function that takes ``Bot, Update, TelegramError`` as
347385
arguments.
348386
"""
349387

@@ -354,7 +392,7 @@ def remove_error_handler(self, callback):
354392
Removes an error handler.
355393
356394
Args:
357-
handler (:obj:`callable`): The error handler to remove.
395+
callback (:obj:`callable`): The error handler to remove.
358396
"""
359397

360398
if callback in self.error_handlers:
@@ -365,7 +403,7 @@ def dispatch_error(self, update, error):
365403
Dispatches an error.
366404
367405
Args:
368-
update (:obj:`str` | :class:`telegram.Update`): The update that caused the error
406+
update (:obj:`str` | :class:`telegram.Update` | None): The update that caused the error
369407
error (:class:`telegram.TelegramError`): The Telegram error that was raised.
370408
"""
371409

‎telegram/utils/promise.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,3 +56,7 @@ def result(self, timeout=None):
5656
if self._exception is not None:
5757
raise self._exception # pylint: disable=raising-bad-type
5858
return self._result
59+
60+
@property
61+
def exception(self):
62+
return self._exception

‎tests/test_updater.py‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -939,6 +939,81 @@ def test_mutualExclusiveTokenBot(self):
939939
def test_noTokenOrBot(self):
940940
self.assertRaises(ValueError, Updater)
941941

942+
def test_dispatcher_handler_flow_continue(self):
943+
passed = []
944+
945+
def start1(b, u):
946+
passed.append('start1')
947+
raise DispatcherHandlerContinue
948+
949+
def start2(b, u):
950+
passed.append('start2')
951+
952+
def start3(b, u):
953+
passed.append('start3')
954+
955+
def error(b, u, e):
956+
passed.append('error')
957+
passed.append(e)
958+
959+
# noinspection PyTypeChecker
960+
update = Update(1, message=Message(1, None, None, None, text='/start', bot=self._bot))
961+
962+
# Without raising Continue everything should work as before
963+
passed = []
964+
dp = Dispatcher(self._bot, Queue())
965+
dp.add_handler(CommandHandler('start', start3))
966+
dp.add_handler(CommandHandler('start', start2))
967+
dp.add_error_handler(error)
968+
dp.process_update(update)
969+
self.assertEqual(passed, ['start3'])
970+
971+
# If Continue raised next handler should be proceed.
972+
passed = []
973+
dp = Dispatcher(self._bot, Queue())
974+
dp.add_handler(CommandHandler('start', start1))
975+
dp.add_handler(CommandHandler('start', start2))
976+
dp.process_update(update)
977+
self.assertEqual(passed, ['start1', 'start2'])
978+
979+
def test_dispatcher_handler_flow_stop(self):
980+
passed = []
981+
982+
def start1(b, u):
983+
passed.append('start1')
984+
raise DispatcherHandlerStop
985+
986+
def start2(b, u):
987+
passed.append('start2')
988+
989+
def start3(b, u):
990+
passed.append('start3')
991+
992+
def error(b, u, e):
993+
passed.append('error')
994+
passed.append(e)
995+
996+
# noinspection PyTypeChecker
997+
update = Update(1, message=Message(1, None, None, None, text='/start', bot=self._bot))
998+
999+
# Without raising Stop everything should work as before
1000+
passed = []
1001+
dp = Dispatcher(self._bot, Queue())
1002+
dp.add_handler(CommandHandler('start', start3), 1)
1003+
dp.add_handler(CommandHandler('start', start2), 2)
1004+
dp.add_error_handler(error)
1005+
dp.process_update(update)
1006+
self.assertEqual(passed, ['start3', 'start2'])
1007+
1008+
# If Stop raised handlers in other groups should not be called.
1009+
passed = []
1010+
dp = Dispatcher(self._bot, Queue())
1011+
dp.add_handler(CommandHandler('start', start1), 1)
1012+
dp.add_handler(CommandHandler('start', start3), 1)
1013+
dp.add_handler(CommandHandler('start', start2), 2)
1014+
dp.process_update(update)
1015+
self.assertEqual(passed, ['start1'])
1016+
9421017

9431018
class MockBot(object):
9441019
def __init__(self,

0 commit comments

Comments
 (0)