fix for python-telegram-bot#123

jh0ker · jh0ker · commit c15a1bc1516d · 2015-12-21T19:36:17.000+01:00
diff --git a/examples/updater_bot.py b/examples/updater_bot.py
@@ -65,12 +65,15 @@ def unknown_command(bot, update):
 
 
 @run_async
-def message(bot, update):
+def message(bot, update, **kwargs):
     """
     Example for an asynchronous handler. It's not guaranteed that replies will
-    be in order when using @run_async.
+    be in order when using @run_async. Also, you have to include **kwargs in
+    your parameter list.
     """
 
+    print(kwargs)
+
     sleep(2)  # IO-heavy operation here
     bot.sendMessage(update.message.chat_id, text='Echo: %s' %
                                                  update.message.text)
diff --git a/telegram/dispatcher.py b/telegram/dispatcher.py
@@ -21,7 +21,9 @@
 
 def run_async(func):
     """
-    Function decorator that will run the function in a new thread.
+    Function decorator that will run the function in a new thread. A function
+    decorated with this will have to include **kwargs in their parameter list,
+    which will contain all optional parameters.
 
     Args:
         func (function): The function to run in the thread.
@@ -31,24 +33,24 @@ def run_async(func):
     """
 
     @wraps(func)
-    def pooled(*args, **kwargs):
+    def pooled(*pargs, **kwargs):
         """
         A wrapper to run a thread in a thread pool
         """
         global running_async, async_lock
-        result = func(*args, **kwargs)
+        result = func(*pargs, **kwargs)
         semaphore.release()
         with async_lock:
             running_async -= 1
         return result
 
     @wraps(func)
-    def async_func(*args, **kwargs):
+    def async_func(*pargs, **kwargs):
         """
         A wrapper to run a function in a thread
         """
         global running_async, async_lock
-        thread = Thread(target=pooled, args=args, kwargs=kwargs)
+        thread = Thread(target=pooled, args=pargs, kwargs=kwargs)
         semaphore.acquire()
         with async_lock:
             running_async += 1
@@ -509,8 +511,7 @@ def dispatchStringRegex(self, update):
 
         Args:
             command (str): The command keyword
-            update (telegram.Update): The Telegram update that contains the
-                command
+            update (str): The string that contains the command
         """
 
         matching_handlers = []
@@ -560,7 +561,7 @@ def dispatchError(self, update, error):
         for handler in self.error_handlers:
             handler(self.bot, update, error)
 
-    def dispatchTo(self, handlers, update):
+    def dispatchTo(self, handlers, update, **kwargs):
         """
         Dispatches an update to a list of handlers.
 
@@ -570,9 +571,9 @@ def dispatchTo(self, handlers, update):
         """
 
         for handler in handlers:
-            self.call_handler(handler, update)
+            self.call_handler(handler, update, **kwargs)
 
-    def call_handler(self, handler, update):
+    def call_handler(self, handler, update, **kwargs):
         """
         Calls an update handler. Checks the handler for keyword arguments and
         fills them, if possible.
@@ -581,20 +582,34 @@ def call_handler(self, handler, update):
             handler (function): An update handler function
             update (any): An update
         """
-        kwargs = {}
+
+        target_kwargs = {}
         fargs = getargspec(handler).args
 
-        if 'update_queue' in fargs:
-            kwargs['update_queue'] = self.update_queue
+        '''
+        async handlers will receive all optional arguments, since we can't
+        their argument list.
+        '''
+
+        is_async = 'pargs' == getargspec(handler).varargs
 
-        if 'args' in fargs:
+        if is_async or 'update_queue' in fargs:
+            target_kwargs['update_queue'] = self.update_queue
+
+        if is_async or 'args' in fargs:
             if isinstance(update, Update):
                 args = update.message.text.split(' ')[1:]
             elif isinstance(update, str):
                 args = update.split(' ')[1:]
             else:
                 args = None
 
-            kwargs['args'] = args
+            target_kwargs['args'] = args
+
+        if is_async or 'groups' in fargs:
+            target_kwargs['groups'] = kwargs.get('groups', None)
+
+        if is_async or 'groupdict' in fargs:
+            target_kwargs['groupdict'] = kwargs.get('groupdict', None)
 
-        handler(self.bot, update, **kwargs)
+        handler(self.bot, update, **target_kwargs)
diff --git a/tests/test_updater.py b/tests/test_updater.py
@@ -79,7 +79,7 @@ def telegramHandlerTest(self, bot, update):
         self.message_count += 1
 
     @run_async
-    def asyncHandlerTest(self, bot, update):
+    def asyncHandlerTest(self, bot, update, **kwargs):
         sleep(1)
         with self.lock:
             self.received_message = update.message.text
@@ -97,6 +97,15 @@ def additionalArgsTest(self, bot, update, update_queue, args):
         elif args[0] == 'noresend':
             pass
 
+    @run_async
+    def asyncAdditionalHandlerTest(self, bot, update, update_queue=None,
+                                   **kwargs):
+        sleep(1)
+        with self.lock:
+            if update_queue is not None:
+                self.received_message = update.message.text
+                self.message_count += 1
+
     def errorRaisingHandlerTest(self, bot, update):
         raise TelegramError(update)
 
@@ -339,6 +348,18 @@ def test_additionalArgs(self):
         self.assertEqual(self.received_message, '/test5 noresend')
         self.assertEqual(self.message_count, 2)
 
+    def test_runAsyncWithAdditionalArgs(self):
+        print('Testing @run_async with additional parameters')
+        bot = MockBot('Test6', messages=2)
+        self.updater.bot = bot
+        d = self.updater.dispatcher
+        d.addTelegramMessageHandler(
+            self.asyncAdditionalHandlerTest)
+        self.updater.start_polling(0.01)
+        sleep(1.2)
+        self.assertEqual(self.received_message, 'Test6')
+        self.assertEqual(self.message_count, 2)
+
     def test_webhook(self):
         print('Testing Webhook')
         bot = MockBot('', messages=0)
