Add has_args Parameter to CommandHandler (#3854)

Co-authored-by: Yao Kuan <chan_yao_kuan@tech.gov.sg>

AUTHORS.rst

@@ -120,6 +120,7 @@ The following wonderful people contributed directly or indirectly to this projec
 - `Wagner Macedo <https://github.com/wagnerluis1982>`_
 - `wjt <https://github.com/wjt>`_
 - `Yaw Danso <https://github.com/dglitxh>`_
+- `Yao Kuan <https://github.com/thatguylah>`_
 - `zeroone2numeral2 <https://github.com/zeroone2numeral2>`_
 - `zeshuaro <https://github.com/zeshuaro>`_
 - `zpavloudis <https://github.com/zpavloudis>`_

telegram/ext/_commandhandler.py

@@ -88,6 +88,14 @@ class CommandHandler(BaseHandler[Update, CCT]):
             :meth:`telegram.ext.Application.process_update`. Defaults to :obj:`True`.
 
             .. seealso:: :wiki:`Concurrency`
+        has_args (:obj:`bool` | :obj:`int`, optional):
+            Determines whether the command handler should process the update or not.
+            If :obj:`True`, the handler will process any non-zero number of args.
+            If :obj:`False`, the handler will only process if there are no args.
+            if :obj:`int`, the handler will only process if there are exactly that many args.
+            Defaults to :obj:`None`, which means the handler will process any or no args.
+
+            .. versionadded:: NEXT.VERSION
 
     Raises:
         :exc:`ValueError`: When the command is too long or has illegal chars.
@@ -100,9 +108,14 @@ class CommandHandler(BaseHandler[Update, CCT]):
         block (:obj:`bool`): Determines whether the return value of the callback should be
             awaited before processing the next handler in
             :meth:`telegram.ext.Application.process_update`.
+        has_args (:obj:`bool` | :obj:`int` | None):
+            Optional argument, otherwise all implementations of :class:`CommandHandler` will break.
+            Defaults to :obj:`None`, which means the handler will process any args or no args.
+
+            .. versionadded:: NEXT.VERSION
     """
 
-    __slots__ = ("commands", "filters")
+    __slots__ = ("commands", "filters", "has_args")
 
     def __init__(
         self,
@@ -110,6 +123,7 @@ class CommandHandler(BaseHandler[Update, CCT]):
         callback: HandlerCallback[Update, CCT, RT],
         filters: Optional[filters_module.BaseFilter] = None,
         block: DVType[bool] = DEFAULT_TRUE,
+        has_args: Optional[Union[bool, int]] = None,
     ):
         super().__init__(callback, block=block)
 
@@ -126,6 +140,28 @@ class CommandHandler(BaseHandler[Update, CCT]):
             filters if filters is not None else filters_module.UpdateType.MESSAGES
         )
 
+        self.has_args: Optional[Union[bool, int]] = has_args
+
+        if (isinstance(self.has_args, int)) and (self.has_args < 0):
+            raise ValueError("CommandHandler argument has_args cannot be a negative integer")
+
+    def _check_correct_args(self, args: List[str]) -> Optional[bool]:
+        """Determines whether the args are correct for this handler. Implemented in check_update().
+        Args:
+            args (:obj:`list`): The args for the handler.
+        Returns:
+            :obj:`bool`: Whether the args are valid for this handler.
+        """
+        # pylint: disable=too-many-boolean-expressions
+        if (
+            (self.has_args is None)
+            or (self.has_args is True and args)
+            or (self.has_args is False and not args)
+            or (isinstance(self.has_args, int) and len(args) == self.has_args)
+        ):
+            return True
+        return False
+
     def check_update(
         self, update: object
     ) -> Optional[Union[bool, Tuple[List[str], Optional[Union[bool, FilterDataDict]]]]]:
@@ -159,6 +195,9 @@ class CommandHandler(BaseHandler[Update, CCT]):
                 ):
                     return None
 
+                if not self._check_correct_args(args):
+                    return None
+
                 filter_result = self.filters.check_update(update)
                 if filter_result:
                     return args, filter_result

tests/ext/test_commandhandler.py

@@ -259,3 +259,32 @@ class TestCommandHandler(BaseTest):
             self.callback_regex2, filters=filters.Regex("one") & filters.Regex("two")
         )
         await self._test_context_args_or_regex(app, handler, command)
+
+    def test_command_has_args(self, bot):
+        """Test CHs with optional has_args specified."""
+        handler_true = CommandHandler(["test"], self.callback_basic, has_args=True)
+        handler_false = CommandHandler(["test"], self.callback_basic, has_args=False)
+        handler_int_one = CommandHandler(["test"], self.callback_basic, has_args=1)
+        handler_int_two = CommandHandler(["test"], self.callback_basic, has_args=2)
+
+        assert is_match(handler_true, make_command_update("/test helloworld", bot=bot))
+        assert not is_match(handler_true, make_command_update("/test", bot=bot))
+
+        assert is_match(handler_false, make_command_update("/test", bot=bot))
+        assert not is_match(handler_false, make_command_update("/test helloworld", bot=bot))
+
+        assert is_match(handler_int_one, make_command_update("/test helloworld", bot=bot))
+        assert not is_match(handler_int_one, make_command_update("/test hello world", bot=bot))
+        assert not is_match(handler_int_one, make_command_update("/test", bot=bot))
+
+        assert is_match(handler_int_two, make_command_update("/test hello world", bot=bot))
+        assert not is_match(handler_int_two, make_command_update("/test helloworld", bot=bot))
+        assert not is_match(handler_int_two, make_command_update("/test", bot=bot))
+
+    def test_command_has_negative_args(self, bot):
+        """Test CHs with optional has_args specified with negative int"""
+        # Assert that CommandHandler cannot be instantiated.
+        with pytest.raises(
+            ValueError, match="CommandHandler argument has_args cannot be a negative integer"
+        ):
+            is_match(CommandHandler(["test"], self.callback_basic, has_args=-1))