add additional arguments 'groups' and 'groupdict' for regex handlers

jh0ker · jh0ker · commit 3c7b037eaaca · 2015-12-21T20:25:31.000+01:00
diff --git a/telegram/dispatcher.py b/telegram/dispatcher.py
@@ -91,6 +91,16 @@ class Dispatcher:
             Example: '/add item1 item2 item3' -> ['item1', 'item2', 'item3']
             For other updates, args will be None
 
+    For regex-based handlers, you can also request information about the match.
+    For all other handlers, these will be None
+
+        groups:
+            A tuple that contains the result of
+            re.match(matcher, ...).groups()
+        groupdict:
+            A dictionary that contains the result of
+            re.match(matcher, ...).groupdict()
+
     Attributes:
 
     Args:
@@ -185,7 +195,7 @@ def processUpdate(self, update):
             self.dispatchStringCommand(update)
             handled = True
         elif type(update) is str:
-            self.dispatchStringRegex(update)
+            self.dispatchRegex(update)
             handled = True
 
         # An error happened while polling
@@ -195,7 +205,7 @@ def processUpdate(self, update):
 
         # Telegram update (regex)
         if isinstance(update, Update):
-            self.dispatchTelegramRegex(update)
+            self.dispatchRegex(update)
             handled = True
 
         # Telegram update (command)
@@ -469,25 +479,30 @@ def dispatchTelegramCommand(self, update):
         else:
             self.dispatchTo(self.unknown_telegram_command_handlers, update)
 
-    def dispatchTelegramRegex(self, update):
+    def dispatchRegex(self, update):
         """
-        Dispatches an update to all regex handlers that match the message
-        string.
+        Dispatches an update to all string or telegram regex handlers that
+        match the string/message content.
 
         Args:
-            command (str): The command keyword
-            update (telegram.Update): The Telegram update that contains the
-                command
+            update (str, Update): The update that should be checked for matches
         """
 
-        matching_handlers = []
-
-        for matcher in self.telegram_regex_handlers:
-            if match(matcher, update.message.text):
-                for handler in self.telegram_regex_handlers[matcher]:
-                    matching_handlers.append(handler)
-
-        self.dispatchTo(matching_handlers, update)
+        if isinstance(update, Update):
+            handlers = self.telegram_regex_handlers
+            to_match = update.message.text
+        elif isinstance(update, str):
+            handlers = self.string_regex_handlers
+            to_match = update
+
+        for matcher in handlers:
+            m = match(matcher, to_match)
+            if m:
+                for handler in handlers[matcher]:
+                    self.call_handler(handler,
+                                      update,
+                                      groups=m.groups(),
+                                      groupdict=m.groupdict())
 
     def dispatchStringCommand(self, update):
         """
@@ -504,25 +519,6 @@ def dispatchStringCommand(self, update):
         else:
             self.dispatchTo(self.unknown_string_command_handlers, update)
 
-    def dispatchStringRegex(self, update):
-        """
-        Dispatches an update to all string regex handlers that match the
-        string.
-
-        Args:
-            command (str): The command keyword
-            update (str): The string that contains the command
-        """
-
-        matching_handlers = []
-
-        for matcher in self.string_regex_handlers:
-            if match(matcher, update):
-                for handler in self.string_regex_handlers[matcher]:
-                    matching_handlers.append(handler)
-
-        self.dispatchTo(matching_handlers, update)
-
     def dispatchType(self, update):
         """
         Dispatches an update of any type.
diff --git a/tests/test_updater.py b/tests/test_updater.py
@@ -89,6 +89,10 @@ def stringHandlerTest(self, bot, update):
         self.received_message = update
         self.message_count += 1
 
+    def regexGroupHandlerTest(self, bot, update, groups=None, groupdict=None):
+        self.received_message = (groups, groupdict)
+        self.message_count += 1
+
     def additionalArgsTest(self, bot, update, update_queue, args):
         self.received_message = update
         self.message_count += 1
@@ -348,6 +352,20 @@ def test_additionalArgs(self):
         self.assertEqual(self.received_message, '/test5 noresend')
         self.assertEqual(self.message_count, 2)
 
+    def test_regexGroupHandler(self):
+        print('Testing optional groups and groupdict parameters')
+        bot = MockBot('', messages=0)
+        self.updater.bot = bot
+        d = self.updater.dispatcher
+        d.addStringRegexHandler('^(This).*?(?P<testgroup>regex group).*',
+                                self.regexGroupHandlerTest)
+        queue = self.updater.start_polling(0.01)
+        queue.put('This is a test message for regex group matching.')
+        sleep(.1)
+        self.assertEqual(self.received_message, (('This', 'regex group'),
+                                                 {'testgroup': 'regex group'}))
+
+
     def test_runAsyncWithAdditionalArgs(self):
         print('Testing @run_async with additional parameters')
         bot = MockBot('Test6', messages=2)
