Merge branch 'more-regex-handlers' of https://github.com/bomjacob/python-telegram-bot into bomjacob-more-regex-handlers

jh0ker · jh0ker · commit 5134f7138040 · 2016-08-06T14:32:05.000+02:00
diff --git a/AUTHORS.rst b/AUTHORS.rst
@@ -13,6 +13,7 @@ The following wonderful people contributed directly or indirectly to this projec
 - `bimmlerd <https://github.com/bimmlerd>`_
 - `ErgoZ Riftbit Vaper <https://github.com/ergoz>`_
 - `franciscod <https://github.com/franciscod>`_
+- `Jacob Bom <https://github.com/bomjacob>`_
 - `JASON0916 <https://github.com/JASON0916>`_
 - `jh0ker <https://github.com/jh0ker>`_
 - `JRoot3D <https://github.com/JRoot3D>`_
diff --git a/telegram/ext/callbackqueryhandler.py b/telegram/ext/callbackqueryhandler.py
@@ -18,14 +18,19 @@
 # along with this program.  If not, see [http://www.gnu.org/licenses/].
 """ This module contains the CallbackQueryHandler class """
 
-from .handler import Handler
+import re
+
+from future.utils import string_types
+
 from telegram import Update
 from telegram.utils.deprecate import deprecate
+from .handler import Handler
 
 
 class CallbackQueryHandler(Handler):
     """
-    Handler class to handle Telegram callback queries.
+    Handler class to handle Telegram callback queries. Optionally based on a regex.
+    Read the documentation of the ``re`` module for more information.
 
     Args:
         callback (function): A function that takes ``bot, update`` as
@@ -39,18 +44,52 @@ class CallbackQueryHandler(Handler):
             ``job_queue`` will be passed to the callback function. It will be a ``JobQueue``
             instance created by the ``Updater`` which can be used to schedule new jobs.
             Default is ``False``.
+        pattern (optional[str or Pattern]): Optional regex pattern. If not ``None`` ``re.match``
+            is used to determine if an update should be handled by this handler.
+        pass_groups (optional[bool]): If the callback should be passed the
+            result of ``re.match(pattern, data).groups()`` as a keyword
+            argument called ``groups``. Default is ``False``
+        pass_groupdict (optional[bool]): If the callback should be passed the
+            result of ``re.match(pattern, data).groupdict()`` as a keyword
+            argument called ``groupdict``. Default is ``False``
     """
 
-    def __init__(self, callback, pass_update_queue=False, pass_job_queue=False):
+    def __init__(self,
+                 callback,
+                 pass_update_queue=False,
+                 pass_job_queue=False,
+                 pattern=None,
+                 pass_groups=False,
+                 pass_groupdict=False):
         super(CallbackQueryHandler, self).__init__(callback,
                                                    pass_update_queue=pass_update_queue,
                                                    pass_job_queue=pass_job_queue)
 
+        if isinstance(pattern, string_types):
+            pattern = re.compile(pattern)
+
+        self.pattern = pattern
+        self.pass_groups = pass_groups
+        self.pass_groupdict = pass_groupdict
+
     def check_update(self, update):
-        return isinstance(update, Update) and update.callback_query
+        if isinstance(update, Update) and update.callback_query:
+            if self.pattern:
+                if update.callback_query.data:
+                    match = re.match(self.pattern, update.callback_query.data)
+                    return bool(match)
+            else:
+                return True
 
     def handle_update(self, update, dispatcher):
         optional_args = self.collect_optional_args(dispatcher)
+        if self.pattern:
+            match = re.match(self.pattern, update.callback_query.data)
+
+            if self.pass_groups:
+                optional_args['groups'] = match.groups()
+            if self.pass_groupdict:
+                optional_args['groupdict'] = match.groupdict()
 
         return self.callback(dispatcher.bot, update, **optional_args)
 
diff --git a/telegram/ext/inlinequeryhandler.py b/telegram/ext/inlinequeryhandler.py
@@ -17,15 +17,19 @@
 # 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 InlineQueryHandler class """
+import re
+
+from future.utils import string_types
 
-from .handler import Handler
 from telegram import Update
 from telegram.utils.deprecate import deprecate
+from .handler import Handler
 
 
 class InlineQueryHandler(Handler):
     """
-    Handler class to handle Telegram inline queries.
+    Handler class to handle Telegram inline queries. Optionally based on a regex. Read the
+    documentation of the ``re`` module for more information.
 
     Args:
         callback (function): A function that takes ``bot, update`` as
@@ -39,18 +43,52 @@ class InlineQueryHandler(Handler):
             ``job_queue`` will be passed to the callback function. It will be a ``JobQueue``
             instance created by the ``Updater`` which can be used to schedule new jobs.
             Default is ``False``.
+        pattern (optional[str or Pattern]): Optional regex pattern. If not ``None`` ``re.match``
+            is used to determine if an update should be handled by this handler.
+        pass_groups (optional[bool]): If the callback should be passed the
+            result of ``re.match(pattern, query).groups()`` as a keyword
+            argument called ``groups``. Default is ``False``
+        pass_groupdict (optional[bool]): If the callback should be passed the
+            result of ``re.match(pattern, query).groupdict()`` as a keyword
+            argument called ``groupdict``. Default is ``False``
     """
 
-    def __init__(self, callback, pass_update_queue=False, pass_job_queue=False):
+    def __init__(self,
+                 callback,
+                 pass_update_queue=False,
+                 pass_job_queue=False,
+                 pattern=None,
+                 pass_groups=False,
+                 pass_groupdict=False):
         super(InlineQueryHandler, self).__init__(callback,
                                                  pass_update_queue=pass_update_queue,
                                                  pass_job_queue=pass_job_queue)
 
+        if isinstance(pattern, string_types):
+            pattern = re.compile(pattern)
+
+        self.pattern = pattern
+        self.pass_groups = pass_groups
+        self.pass_groupdict = pass_groupdict
+
     def check_update(self, update):
-        return isinstance(update, Update) and update.inline_query
+        if isinstance(update, Update) and update.inline_query:
+            if self.pattern:
+                if update.inline_query.query:
+                    match = re.match(self.pattern, update.inline_query.query)
+                    return bool(match)
+            else:
+                return True
 
     def handle_update(self, update, dispatcher):
         optional_args = self.collect_optional_args(dispatcher)
+        if self.pattern:
+            match = re.match(self.pattern, update.inline_query.query)
+
+            if self.pass_groups:
+                optional_args['groups'] = match.groups()
+            if self.pass_groupdict:
+                optional_args['groupdict'] = match.groupdict()
 
         return self.callback(dispatcher.bot, update, **optional_args)
 
