Added prefix='/' : str argument to CommandHandler

JosXa · JosXa · commit b981ea9b4baa · 2018-02-25T22:40:19.000+01:00
diff --git a/telegram/ext/commandhandler.py b/telegram/ext/commandhandler.py
@@ -21,8 +21,8 @@
 
 from future.utils import string_types
 
-from .handler import Handler
 from telegram import Update
+from .handler import Handler
 
 
 class CommandHandler(Handler):
@@ -39,6 +39,8 @@ class CommandHandler(Handler):
             Filters.
         allow_edited (:obj:`bool`): Optional. Determines Whether the handler should also accept
             edited messages.
+        prefix (:obj:`bool`): Optional. Denotes the leading character of commands with this
+            handler.
         pass_args (:obj:`bool`): Optional. Determines whether the handler should be passed
             ``args``.
         pass_update_queue (:obj:`bool`): Optional. Determines whether ``update_queue`` will be
@@ -68,6 +70,8 @@ class CommandHandler(Handler):
             operators (& for and, | for or, ~ for not).
         allow_edited (:obj:`bool`, optional): Determines whether the handler should also accept
             edited messages. Default is ``False``.
+        prefix (:obj:`str`, optional): Denotes the leading character to commands with this handler.
+            Must be exactly one character or an empty string. Default is ``/``.
         pass_args (:obj:`bool`, optional): Determines whether the handler should be passed the
             arguments passed to the command as a keyword argument called ``args``. It will contain
             a list of strings, which is the text following the command split on single or
@@ -92,6 +96,7 @@ def __init__(self,
                  callback,
                  filters=None,
                  allow_edited=False,
+                 prefix='/',
                  pass_args=False,
                  pass_update_queue=False,
                  pass_job_queue=False,
@@ -110,6 +115,15 @@ def __init__(self,
             self.command = [x.lower() for x in command]
         self.filters = filters
         self.allow_edited = allow_edited
+        if prefix is None or not isinstance(prefix, string_types) or len(prefix) > 1:
+            raise ValueError("The prefix argument to CommandHandler must be a single character or "
+                             "an empty string.")
+        if prefix == ' ':
+            # A leading space character is stripped away in Telegram messages, so we can consider
+            # it as an empty string prefix.
+            self.prefix = ''
+        else:
+            self.prefix = prefix
         self.pass_args = pass_args
 
         # We put this up here instead of with the rest of checking code
@@ -133,10 +147,16 @@ def check_update(self, update):
                 and (update.message or update.edited_message and self.allow_edited)):
             message = update.message or update.edited_message
 
-            if message.text and message.text.startswith('/') and len(message.text) > 1:
-                command = message.text[1:].split(None, 1)[0].split('@')
-                command.append(
-                    message.bot.username)  # in case the command was send without a username
+            empty_leader = self.prefix == ''
+            command_text_nonempty = message.text and len(message.text) > 1
+            is_valid_command = command_text_nonempty and (empty_leader or
+                                                          message.text.startswith(self.prefix))
+            if is_valid_command:
+                leader_stripped = message.text if empty_leader else message.text[1:]
+                command = leader_stripped.split(None, 1)[0].split('@')
+
+                # Append the bot's username in any case
+                command.append(message.bot.username)
 
                 if self.filters is None:
                     res = True
diff --git a/tests/test_commandhandler.py b/tests/test_commandhandler.py
@@ -51,6 +51,11 @@ def message(bot):
     return Message(1, None, None, None, bot=bot)
 
 
+@pytest.fixture(scope='function')
+def prefixes():
+    return ['/', '#', '', ' ']
+
+
 class TestCommandHandler(object):
     test_flag = False
 
@@ -75,62 +80,72 @@ def callback_queue_1(self, bot, update, job_queue=None, update_queue=None):
     def callback_queue_2(self, bot, update, job_queue=None, update_queue=None):
         self.test_flag = (job_queue is not None) and (update_queue is not None)
 
-    def ch_callback_args(self, bot, update, args):
-        if update.message.text == '/test':
+    def ch_callback_args(self, bot, update, args, prefix):
+        if update.message.text == '{}test'.format(prefix):
             self.test_flag = len(args) == 0
-        elif update.message.text == '/test@{}'.format(bot.username):
+        elif update.message.text == '{}test@{}'.format(prefix, bot.username):
             self.test_flag = len(args) == 0
         else:
             self.test_flag = args == ['one', 'two']
 
-    def test_basic(self, dp, message):
-        handler = CommandHandler('test', self.callback_basic)
-        dp.add_handler(handler)
+    def test_basic(self, dp, message, prefixes):
+        for p in prefixes:
+            handler = CommandHandler('test', self.callback_basic, prefix=p)
+            dp.add_handler(handler)
 
-        message.text = '/test'
-        assert handler.check_update(Update(0, message))
-        dp.process_update(Update(0, message))
-        assert self.test_flag
+            message.text = '{}test'.format(p)
+            assert handler.check_update(Update(0, message))
+            dp.process_update(Update(0, message))
+            assert self.test_flag
 
-        message.text = '/nottest'
-        assert not handler.check_update(Update(0, message))
+            message.text = '{}nottest'.format(p)
+            assert not handler.check_update(Update(0, message))
 
-        message.text = 'test'
-        assert not handler.check_update(Update(0, message))
+            message.text = 'test'
+            result = handler.check_update(Update(0, message))
+            if p in ['', ' ']:
+                assert result
+            else:
+                assert not result
 
-        message.text = 'not /test at start'
-        assert not handler.check_update(Update(0, message))
+            message.text = 'not {}test at start'.format(p)
+            assert not handler.check_update(Update(0, message))
 
-    def test_command_list(self, message):
-        handler = CommandHandler(['test', 'start'], self.callback_basic)
+            dp.remove_handler(handler)
 
-        message.text = '/test'
-        assert handler.check_update(Update(0, message))
+    def test_command_list(self, message, prefixes):
+        for p in prefixes:
+            handler = CommandHandler(['test', 'start'], self.callback_basic, prefix=p)
 
-        message.text = '/start'
-        assert handler.check_update(Update(0, message))
+            message.text = '{}test'.format(p)
+            assert handler.check_update(Update(0, message))
 
-        message.text = '/stop'
-        assert not handler.check_update(Update(0, message))
+            message.text = '{}start'.format(p)
+            assert handler.check_update(Update(0, message))
 
-    def test_edited(self, message):
-        handler = CommandHandler('test', self.callback_basic, allow_edited=False)
+            message.text = '{}stop'.format(p)
+            assert not handler.check_update(Update(0, message))
 
-        message.text = '/test'
-        assert handler.check_update(Update(0, message))
-        assert not handler.check_update(Update(0, edited_message=message))
-        handler.allow_edited = True
-        assert handler.check_update(Update(0, message))
-        assert handler.check_update(Update(0, edited_message=message))
+    def test_edited(self, message, prefixes):
+        for p in prefixes:
+            handler = CommandHandler('test', self.callback_basic, allow_edited=False, prefix=p)
 
-    def test_directed_commands(self, message):
-        handler = CommandHandler('test', self.callback_basic)
+            message.text = '{}test'.format(p)
+            assert handler.check_update(Update(0, message))
+            assert not handler.check_update(Update(0, edited_message=message))
+            handler.allow_edited = True
+            assert handler.check_update(Update(0, message))
+            assert handler.check_update(Update(0, edited_message=message))
 
-        message.text = '/test@{}'.format(message.bot.username)
-        assert handler.check_update(Update(0, message))
+    def test_directed_commands(self, message, prefixes):
+        for p in prefixes:
+            handler = CommandHandler('test', self.callback_basic, prefix=p)
 
-        message.text = '/test@otherbot'
-        assert not handler.check_update(Update(0, message))
+            message.text = '{}test@{}'.format(p, message.bot.username)
+            assert handler.check_update(Update(0, message))
+
+            message.text = '{}test@otherbot'.format(p)
+            assert not handler.check_update(Update(0, message))
 
     def test_with_filter(self, message):
         handler = CommandHandler('test', self.callback_basic, Filters.group)
@@ -142,53 +157,74 @@ def test_with_filter(self, message):
         message.chat = Chat(23, 'private')
         assert not handler.check_update(Update(0, message))
 
-    def test_pass_args(self, dp, message):
-        handler = CommandHandler('test', self.ch_callback_args, pass_args=True)
-        dp.add_handler(handler)
-
-        message.text = '/test'
-        dp.process_update(Update(0, message=message))
-        assert self.test_flag
-
-        self.test_flag = False
-        message.text = '/test@{}'.format(message.bot.username)
-        dp.process_update(Update(0, message=message))
-        assert self.test_flag
-
-        self.test_flag = False
-        message.text = '/test one two'
-        dp.process_update(Update(0, message=message))
-        assert self.test_flag
-
-        self.test_flag = False
-        message.text = '/test@{} one two'.format(message.bot.username)
-        dp.process_update(Update(0, message=message))
-        assert self.test_flag
-
-    def test_newline(self, dp, message):
-        handler = CommandHandler('test', self.callback_basic)
-        dp.add_handler(handler)
+    def test_pass_args(self, dp, message, prefixes):
+        for p in prefixes:
+            handler = CommandHandler(
+                'test',
+                lambda bot, update, args: self.ch_callback_args(bot, update, args, p),
+                pass_args=True,
+                prefix=p
+            )
+            dp.add_handler(handler)
+
+            self.test_flag = False
+            message.text = '{}test'.format(p)
+            dp.process_update(Update(0, message=message))
+            assert self.test_flag
+
+            self.test_flag = False
+            message.text = '{}test@{}'.format(p, message.bot.username)
+            dp.process_update(Update(0, message=message))
+            assert self.test_flag
+
+            self.test_flag = False
+            message.text = '{}test one two'.format(p)
+            dp.process_update(Update(0, message=message))
+            assert self.test_flag
+
+            self.test_flag = False
+            message.text = '{}test@{} one two'.format(p, message.bot.username)
+            dp.process_update(Update(0, message=message))
+            assert self.test_flag
+
+            dp.remove_handler(handler)
+
+    def test_newline(self, dp, message, prefixes):
+        for p in prefixes:
+            handler = CommandHandler('test', self.callback_basic, prefix=p)
+            dp.add_handler(handler)
+
+            message.text = '{}test\nfoobar'.format(p)
+            assert handler.check_update(Update(0, message))
+            dp.process_update(Update(0, message))
+            assert self.test_flag
+
+            dp.remove_handler(handler)
+
+    def test_single_char(self, message):
+        # Regression test for
+        # https://github.com/python-telegram-bot/python-telegram-bot/issues/871
 
-        message.text = '/test\nfoobar'
-        assert handler.check_update(Update(0, message))
-        dp.process_update(Update(0, message))
-        assert self.test_flag
+        message.text = 'a'
 
-    def test_single_char(self, dp, message):
-        # Regression test for https://github.com/python-telegram-bot/python-telegram-bot/issues/871
-        handler = CommandHandler('test', self.callback_basic)
-        dp.add_handler(handler)
+        normal_handler = CommandHandler('test', self.callback_basic)
+        assert not normal_handler.check_update(Update(0, message))
 
-        message.text = 'a'
-        assert not handler.check_update(Update(0, message))
+        empty_prefix_handler = CommandHandler('test', self.callback_basic, prefix='')
+        assert not empty_prefix_handler.check_update(Update(0, message))
 
-    def test_single_slash(self, dp, message):
-        # Regression test for https://github.com/python-telegram-bot/python-telegram-bot/issues/871
-        handler = CommandHandler('test', self.callback_basic)
-        dp.add_handler(handler)
+    def test_single_prefix(self, dp, message, prefixes):
+        for p in prefixes:
+            # Regression test for
+            # https://github.com/python-telegram-bot/python-telegram-bot/issues/871
+            handler = CommandHandler('test', self.callback_basic, prefix=p)
+            dp.add_handler(handler)
 
-        message.text = '/'
-        assert not handler.check_update(Update(0, message))
+            # Note: In praxis, it is not possible to send empty messages.
+            #       We will test this case nonetheless
+            message.text = p
+            assert not handler.check_update(Update(0, message))
+            dp.remove_handler(handler)
 
     def test_pass_user_or_chat_data(self, dp, message):
         handler = CommandHandler('test', self.callback_data_1, pass_user_data=True)
