Merge pull request #157 from jlmadurga/context-for-handlers

Add context to dispatcher

AUTHORS.rst

@@ -16,6 +16,7 @@ The following wonderful people contributed directly or indirectly to this projec
 - `JASON0916 <https://github.com/JASON0916>`_
 - `jh0ker <https://github.com/jh0ker>`_
 - `JRoot3D <https://github.com/JRoot3D>`_
+- `jlmadurga <https://github.com/jlmadurga>`_
 - `macrojames <https://github.com/macrojames>`_
 - `naveenvhegde <https://github.com/naveenvhegde>`_
 - `njittam <https://github.com/njittam>`_

telegram/__init__.py

@@ -47,8 +47,10 @@ from .update import Update
 from .bot import Bot
 from .dispatcher import Dispatcher
 from .jobqueue import JobQueue
+from .updatequeue import UpdateQueue
 from .updater import Updater
 
+
 __author__ = 'devs@python-telegram-bot.org'
 __version__ = '3.2.0'
 __all__ = ('Bot', 'Updater', 'Dispatcher', 'Emoji', 'TelegramError',
@@ -56,4 +58,5 @@ __all__ = ('Bot', 'Updater', 'Dispatcher', 'Emoji', 'TelegramError',
            'ReplyKeyboardMarkup', 'UserProfilePhotos', 'ChatAction',
            'Location', 'Contact', 'Video', 'Sticker', 'Document', 'File',
            'Audio', 'PhotoSize', 'Chat', 'Update', 'ParseMode', 'Message',
-           'User', 'TelegramObject', 'NullHandler', 'Voice', 'JobQueue')
+           'User', 'TelegramObject', 'NullHandler', 'Voice', 'JobQueue',
+           'UpdateQueue')

telegram/dispatcher.py

@@ -108,6 +108,13 @@ class Dispatcher:
             Example: '/add item1 item2 item3' -> ['item1', 'item2', 'item3']
             For other updates, args will be None
 
+    In some cases handlers may need some context data to process the update. To
+    procedure just queue in  update_queue.put(update, context=context) or
+    processUpdate(update,context=context).
+
+        context:
+            Extra data for handling updates.
+
     For regex-based handlers, you can also request information about the match.
     For all other handlers, these will be None
 
@@ -121,7 +128,7 @@ class Dispatcher:
     Args:
         bot (telegram.Bot): The bot object that should be passed to the
             handlers
-        update_queue (queue.Queue): The synchronized queue that will
+        update_queue (UpdateQueue): The synchronized queue that will
             contain the updates.
     """
     def __init__(self, bot, update_queue, workers=4):
@@ -171,14 +178,15 @@ class Dispatcher:
                 try:
                     # Pop update from update queue.
                     # Blocks if no updates are available.
-                    update = self.update_queue.get()
+                    update, context = self.update_queue.get(context=True)
 
                     if type(update) is self._Stop:
                         self.running = False
                         break
 
-                    self.processUpdate(update)
+                    self.processUpdate(update, context)
-                    self.logger.debug('Processed Update: %s' % update)
+                    self.logger.debug('Processed Update: %s with context %s'
+                                      % (update, context))
 
                 # Dispatch any errors
                 except TelegramError as te:
@@ -204,7 +212,7 @@ class Dispatcher:
                 while self.running:
                     sleep(0.1)
 
-    def processUpdate(self, update):
+    def processUpdate(self, update, context=None):
         """
         Processes a single update.
 
@@ -217,15 +225,15 @@ class Dispatcher:
         # Custom type handlers
         for t in self.type_handlers:
             if isinstance(update, t):
-                self.dispatchType(update)
+                self.dispatchType(update, context)
                 handled = True
 
         # string update
         if type(update) is str and update.startswith('/'):
-            self.dispatchStringCommand(update)
+            self.dispatchStringCommand(update, context)
             handled = True
         elif type(update) is str:
-            self.dispatchRegex(update)
+            self.dispatchRegex(update, context)
             handled = True
 
         # An error happened while polling
@@ -235,18 +243,18 @@ class Dispatcher:
 
         # Telegram update (regex)
         if isinstance(update, Update):
-            self.dispatchRegex(update)
+            self.dispatchRegex(update, context)
             handled = True
 
         # Telegram update (command)
         if isinstance(update, Update) \
                 and update.message.text.startswith('/'):
-            self.dispatchTelegramCommand(update)
+            self.dispatchTelegramCommand(update, context)
             handled = True
 
         # Telegram update (message)
         elif isinstance(update, Update):
-            self.dispatchTelegramMessage(update)
+            self.dispatchTelegramMessage(update, context)
             handled = True
 
         # Update not recognized
@@ -492,7 +500,7 @@ class Dispatcher:
                 and handler in self.type_handlers[the_type]:
             self.type_handlers[the_type].remove(handler)
 
-    def dispatchTelegramCommand(self, update):
+    def dispatchTelegramCommand(self, update, context=None):
         """
         Dispatches an update that contains a command.
 
@@ -505,11 +513,13 @@ class Dispatcher:
         command = update.message.text.split(' ')[0][1:].split('@')[0]
 
         if command in self.telegram_command_handlers:
-            self.dispatchTo(self.telegram_command_handlers[command], update)
+            self.dispatchTo(self.telegram_command_handlers[command], update,
+                            context=context)
         else:
-            self.dispatchTo(self.unknown_telegram_command_handlers, update)
+            self.dispatchTo(self.unknown_telegram_command_handlers, update,
+                            context=context)
 
-    def dispatchRegex(self, update):
+    def dispatchRegex(self, update, context=None):
         """
         Dispatches an update to all string or telegram regex handlers that
         match the string/message content.
@@ -532,9 +542,10 @@ class Dispatcher:
                     self.call_handler(handler,
                                       update,
                                       groups=m.groups(),
-                                      groupdict=m.groupdict())
+                                      groupdict=m.groupdict(),
+                                      context=context)
 
-    def dispatchStringCommand(self, update):
+    def dispatchStringCommand(self, update, context=None):
         """
         Dispatches a string-update that contains a command.
 
@@ -545,11 +556,13 @@ class Dispatcher:
         command = update.split(' ')[0][1:]
 
         if command in self.string_command_handlers:
-            self.dispatchTo(self.string_command_handlers[command], update)
+            self.dispatchTo(self.string_command_handlers[command], update,
+                            context=context)
         else:
-            self.dispatchTo(self.unknown_string_command_handlers, update)
+            self.dispatchTo(self.unknown_string_command_handlers, update,
+                            context=context)
 
-    def dispatchType(self, update):
+    def dispatchType(self, update, context=None):
         """
         Dispatches an update of any type.
 
@@ -559,9 +572,9 @@ class Dispatcher:
 
         for t in self.type_handlers:
             if isinstance(update, t):
-                self.dispatchTo(self.type_handlers[t], update)
+                self.dispatchTo(self.type_handlers[t], update, context=context)
 
-    def dispatchTelegramMessage(self, update):
+    def dispatchTelegramMessage(self, update, context=None):
         """
         Dispatches an update that contains a regular message.
 
@@ -570,7 +583,8 @@ class Dispatcher:
                 message.
         """
 
-        self.dispatchTo(self.telegram_message_handlers, update)
+        self.dispatchTo(self.telegram_message_handlers, update,
+                        context=context)
 
     def dispatchError(self, update, error):
         """
@@ -635,4 +649,7 @@ class Dispatcher:
         if is_async or 'groupdict' in fargs:
             target_kwargs['groupdict'] = kwargs.get('groupdict', None)
 
+        if is_async or 'context' in fargs:
+            target_kwargs['context'] = kwargs.get('context', None)
+
         handler(self.bot, update, **target_kwargs)

telegram/updatequeue.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python
+#
+# A library that provides a Python interface to the Telegram Bot API
+# Copyright (C) 2015-2016
+# Leandro Toledo de Souza <devs@python-telegram-bot.org>
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program 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 Public License for more details.
+#
+# You should have received a copy of the GNU Lesser Public License
+# along with this program.  If not, see [http://www.gnu.org/licenses/].
+
+"""This module contains the class UpdateQueue to override standard
+Queue."""
+
+
+# Adjust for differences in Python versions
+try:
+    from queue import Queue
+except ImportError:
+    from Queue import Queue
+
+
+class UpdateQueue(Queue):
+    """
+    This class overrides standard Queues. Allows you to de/queue context
+    data apart from the handled `update`
+    """
+
+    def put(self, item, block=True, timeout=None, context=None):
+        """
+        Put an item into the queue with context data if provided as a
+        tuple (item, context). Overrides standard Queue.put method.
+
+        Args:
+            update (any): handled by the dispatcher
+            context (any): extra data to use in handlers
+        """
+        Queue.put(self, (item, context), block, timeout)
+
+    def get(self, block=True, timeout=None, context=False):
+        """
+        Remove and return an item from the queue. A tuple of
+        (update, context) if requested. Overrides standard Queue.get
+        method.
+
+        Args:
+            context (boolean): set true to get (update, context)
+        """
+        if not context:
+            return Queue.get(self, block, timeout)[0]
+        return Queue.get(self, block, timeout)

telegram/updater.py

@@ -29,15 +29,9 @@ from time import sleep
 import subprocess
 from signal import signal, SIGINT, SIGTERM, SIGABRT
 from telegram import (Bot, TelegramError, dispatcher, Dispatcher,
-                      NullHandler, JobQueue)
+                      NullHandler, JobQueue, UpdateQueue)
 from telegram.utils.webhookhandler import (WebhookServer, WebhookHandler)
 
-# Adjust for differences in Python versions
-try:
-    from Queue import Queue
-except ImportError:
-    from queue import Queue
-
 try:
     from urllib2 import URLError
 except ImportError:
@@ -89,7 +83,7 @@ class Updater:
             self.bot = bot
         else:
             self.bot = Bot(token, base_url)
-        self.update_queue = Queue()
+        self.update_queue = UpdateQueue()
         self.job_queue = JobQueue(self.bot, job_queue_tick_interval)
         self.dispatcher = Dispatcher(self.bot, self.update_queue,
                                      workers=workers)

tests/test_updater.py

@@ -109,6 +109,11 @@ class UpdaterTest(BaseTest, unittest.TestCase):
             update_queue.put('/test5 noresend')
         elif args[0] == 'noresend':
             pass
+        
+    def contextTest(self, bot, update, context):
+        self.received_message = update
+        self.message_count += 1
+        self.context = context
 
     @run_async
     def asyncAdditionalHandlerTest(self, bot, update, update_queue=None,
@@ -348,6 +353,20 @@ class UpdaterTest(BaseTest, unittest.TestCase):
         sleep(.1)
         self.assertEqual(self.received_message, '/test5 noresend')
         self.assertEqual(self.message_count, 2)
+        
+    def test_context(self):
+        print('Testing context for handlers')
+        context = "context_data"
+        self._setup_updater('', messages=0)
+        self.updater.dispatcher.addStringCommandHandler(
+            'test_context', self.contextTest)
+        
+        queue = self.updater.start_polling(0.01)
+        queue.put('/test_context', context=context)
+        sleep(.5)
+        self.assertEqual(self.received_message, '/test_context')
+        self.assertEqual(self.message_count, 1)
+        self.assertEqual(self.context, context)
 
     def test_regexGroupHandler(self):
         print('Testing optional groups and groupdict parameters')