Improved CommandHandler (python-telegram-bot#613)

evgfilim1 · jh0ker · commit 33512ffd2e85 · 2017-05-19T20:21:37.000+02:00
* Improved CommandHandler

Now you can pass list of commands instead of one command

* Added tests for list of commands

* Return backward compatibility

Renamed `commands` to `command` in CommandHandler

* Added test for a command not in the list

* Fixed py2 unicode command support in `CommandHandler`
diff --git a/telegram/ext/commandhandler.py b/telegram/ext/commandhandler.py
@@ -30,7 +30,8 @@ class CommandHandler(Handler):
     name and/or some additional text.
 
     Args:
-        command (str): The name of the command this handler should listen for.
+        command (str|list): The name of the command or list of command this handler should
+            listen for.
         callback (function): A function that takes ``bot, update`` as
             positional arguments. It will be called when the ``check_update``
             has determined that an update should be processed by this handler.
@@ -79,7 +80,15 @@ def __init__(self,
             pass_job_queue=pass_job_queue,
             pass_user_data=pass_user_data,
             pass_chat_data=pass_chat_data)
-        self.command = command
+        try:
+            _str = basestring  # Python 2
+        except NameError:
+            _str = str  # Python 3
+
+        if isinstance(command, _str):
+            self.command = [command]
+        else:
+            self.command = command
         self.filters = filters
         self.allow_edited = allow_edited
         self.pass_args = pass_args
@@ -108,7 +117,7 @@ def check_update(self, update):
                 else:
                     res = self.filters(message)
 
-                return res and (message.text.startswith('/') and command[0] == self.command
+                return res and (message.text.startswith('/') and command[0] in self.command
                                 and command[1].lower() == message.bot.username.lower())
             else:
                 return False
diff --git a/tests/test_updater.py b/tests/test_updater.py
@@ -300,6 +300,35 @@ def test_filterNotPassTelegramCommandHandler(self):
         sleep(.1)
         self.assertTrue(None is self.received_message)
 
+    def test_CommandHandler_commandList(self):
+        self._setup_updater('', messages=0)
+        handler = CommandHandler(['foo', 'bar', 'spameggs'], self.telegramHandlerTest)
+        self.updater.dispatcher.add_handler(handler)
+        bot = self.updater.bot
+        user = User(0, 'TestUser')
+        queue = self.updater.start_polling(0.01)
+
+        message = Message(0, user, 0, None, text='/foo', bot=bot)
+        queue.put(Update(0, message=message))
+        sleep(.1)
+        self.assertEqual(self.received_message, '/foo')
+
+        message.text = '/bar'
+        queue.put(Update(1, message=message))
+        sleep(.1)
+        self.assertEqual(self.received_message, '/bar')
+
+        message.text = '/spameggs'
+        queue.put(Update(2, message=message))
+        sleep(.1)
+        self.assertEqual(self.received_message, '/spameggs')
+
+        self.reset()
+        message.text = '/not_in_list'
+        queue.put(Update(3, message=message))
+        sleep(.1)
+        self.assertTrue(self.received_message is None)
+
     def test_addRemoveStringRegexHandler(self):
         self._setup_updater('', messages=0)
         d = self.updater.dispatcher
