Ruggedize a few types

Author: tleonhardt

cmd2/annotated.py

@@ -290,6 +290,7 @@ def do_build(self, target: str, common: CommonArgs):
     TypeVar,
     Union,
     Unpack,
+    cast,
     get_args,
     get_origin,
     get_type_hints,
@@ -321,7 +322,7 @@ def do_build(self, target: str, common: CommonArgs):
 _NargsValue = int | str | tuple[int] | tuple[int, int] | tuple[int, float]
 
 # Parser construction works with a command method before or after descriptor binding.
-_CommandFunc: TypeAlias = BoundCommandFunc | UnboundCommandFunc[Any, ...]
+_CommandFunc: TypeAlias = BoundCommandFunc[...] | UnboundCommandFunc[Any, ...]
 
 
 class Cmd2ParserKwargs(TypedDict, total=False):
@@ -2880,28 +2881,36 @@ def parser_builder() -> Cmd2ArgumentParser:
 _DecoratorParams = ParamSpec("_DecoratorParams")
 
 
-class _AnnotatedCommand(Protocol[_CommandParams]):
+class _AnnotatedCommand(Protocol[CmdOrSetT, _CommandParams]):
     """A callable command method with the metadata used by ``with_annotated``."""
 
     __name__: str
     __qualname__: str
 
-    def __call__(self, *args: _CommandParams.args, **kwargs: _CommandParams.kwargs) -> bool | None: ...
+    def __call__(self, __self: CmdOrSetT, /, *args: _CommandParams.args, **kwargs: _CommandParams.kwargs) -> bool | None: ...
+
+    @overload
+    def __get__(self, instance: None, owner: Any = ...) -> "_AnnotatedCommand[CmdOrSetT, _CommandParams]": ...
+
+    @overload
+    def __get__(self, instance: CmdOrSetT, owner: Any = ...) -> BoundCommandFunc[_CommandParams]: ...
 
 
 class _WithAnnotatedDecorator(Protocol):
     """The signature-preserving decorator ``with_annotated(...)`` returns (generic per call)."""
 
-    def __call__(self, fn: _AnnotatedCommand[_CommandParams], /) -> _AnnotatedCommand[_CommandParams]: ...
+    def __call__(
+        self, fn: _AnnotatedCommand[CmdOrSetT, _CommandParams], /
+    ) -> _AnnotatedCommand[CmdOrSetT, _CommandParams]: ...
 
 
 def _build_subcommand_handler(
-    func: _AnnotatedCommand[_CommandParams],
+    func: _AnnotatedCommand[Any, ...],
     subcommand_to: str,
     *,
     base_command: bool = False,
     options: _ParserBuildOptions,
-) -> tuple[_AnnotatedCommand[_CommandParams], str, Callable[[], Cmd2ArgumentParser]]:
+) -> tuple[_AnnotatedCommand[Any, ...], str, Callable[[], Cmd2ArgumentParser]]:
     """Build a subcommand's parser and a handler that unpacks the Namespace into typed kwargs.
 
     :param func: the subcommand handler function
@@ -2927,7 +2936,7 @@ def _build_subcommand_handler(
     _leading_names, _var_positional_name = _var_positional_call_plan(func)
 
     @functools.wraps(func)
-    def handler(*args: _CommandParams.args, **kwargs: _CommandParams.kwargs) -> bool | None:
+    def handler(*args: Any, **kwargs: Any) -> bool | None:
         """Unpack Namespace into typed kwargs for the subcommand handler."""
         if kwargs:
             raise TypeError("subcommand handlers do not accept keyword arguments")
@@ -2951,11 +2960,11 @@ def handler(*args: _CommandParams.args, **kwargs: _CommandParams.kwargs) -> bool
         raise TypeError("annotated command functions must return bool or None")
 
     parser_builder = _make_parser_builder(func, skip_params=_SKIP_PARAMS, base_command=base_command, options=options)
-    return handler, subcmd_name, parser_builder
+    return cast(_AnnotatedCommand[Any, ...], handler), subcmd_name, parser_builder
 
 
 @overload
-def with_annotated(func: _AnnotatedCommand[_CommandParams], /) -> _AnnotatedCommand[_CommandParams]: ...
+def with_annotated(func: _AnnotatedCommand[CmdOrSetT, _CommandParams], /) -> _AnnotatedCommand[CmdOrSetT, _CommandParams]: ...
 
 
 @overload
@@ -2982,7 +2991,7 @@ def with_annotated(
 
 
 def with_annotated(
-    func: _AnnotatedCommand[_CommandParams] | None = None,
+    func: _AnnotatedCommand[CmdOrSetT, _CommandParams] | None = None,
     *,
     ns_provider: Callable[..., argparse.Namespace] | None = None,
     preserve_quotes: bool = False,
@@ -3000,7 +3009,7 @@ def with_annotated(
     subcommand_title: str | None = None,
     subcommand_description: str | None = None,
     **parser_kwargs: Unpack[Cmd2ParserKwargs],
-) -> _AnnotatedCommand[_CommandParams] | _WithAnnotatedDecorator:
+) -> _AnnotatedCommand[CmdOrSetT, _CommandParams] | _WithAnnotatedDecorator:
     """Decorate a ``do_*`` method to build its argparse parser from type annotations.
 
     :param func: the command function (when used without parentheses)
@@ -3062,8 +3071,8 @@ def with_annotated(
     )
 
     def decorator(
-        command_func: _AnnotatedCommand[_DecoratorParams],
-    ) -> _AnnotatedCommand[_DecoratorParams]:
+        command_func: _AnnotatedCommand[CmdOrSetT, _DecoratorParams],
+    ) -> _AnnotatedCommand[CmdOrSetT, _DecoratorParams]:
         if with_unknown_args:
             unknown_param = inspect.signature(command_func).parameters.get("_unknown")
             if unknown_param is None:
@@ -3093,7 +3102,7 @@ def decorator(
                 parser_source=subcmd_parser_builder,
             )
             setattr(handler, constants.SUBCOMMAND_ATTR_SPEC, subcommand_spec)
-            return handler
+            return cast(_AnnotatedCommand[CmdOrSetT, _DecoratorParams], handler)
 
         command_name = command_func.__name__[len(constants.COMMAND_FUNC_PREFIX) :]
 
@@ -3170,9 +3179,9 @@ def cmd_wrapper(*args: Any, **kwargs: Any) -> bool | None:
         )
         setattr(cmd_wrapper, constants.ARGPARSE_COMMAND_ATTR_SPEC, argparse_command_spec)
 
-        return cmd_wrapper
+        return cast(_AnnotatedCommand[CmdOrSetT, _DecoratorParams], cmd_wrapper)
 
     # Support both @with_annotated and @with_annotated(...)
     if func is not None:
         return decorator(func)
-    return decorator
+    return cast(_WithAnnotatedDecorator, decorator)

cmd2/cmd2.py

@@ -228,7 +228,7 @@ class DisabledCommand(NamedTuple):
     This data is used to restore its functions when the command is enabled.
     """
 
-    command_func: BoundCommandFunc
+    command_func: BoundCommandFunc[...]
     help_func: Callable[[], Any] | None
     completer_func: BoundCompleter | None
 
@@ -251,14 +251,14 @@ def __init__(self, cmd_app: "Cmd") -> None:
         self._parsers: dict[str, Cmd2ArgumentParser] = {}
 
     @staticmethod
-    def _fully_qualified_name(command_method: BoundCommandFunc) -> str:
+    def _fully_qualified_name(command_method: BoundCommandFunc[...]) -> str:
         """Return the fully qualified name of a method or None if a method wasn't passed in."""
         try:
             return f"{command_method.__module__}.{command_method.__qualname__}"
         except AttributeError:
             return ""
 
-    def __contains__(self, command_method: BoundCommandFunc) -> bool:
+    def __contains__(self, command_method: BoundCommandFunc[...]) -> bool:
         """Return whether a given method's parser is in self.
 
         If the parser does not yet exist, it will be created if applicable.
@@ -267,7 +267,7 @@ def __contains__(self, command_method: BoundCommandFunc) -> bool:
         parser = self.get(command_method)
         return bool(parser)
 
-    def get(self, command_method: BoundCommandFunc) -> Cmd2ArgumentParser | None:
+    def get(self, command_method: BoundCommandFunc[...]) -> Cmd2ArgumentParser | None:
         """Return a given method's parser or None if the method is not argparse-based.
 
         If the parser does not yet exist, it will be created.
@@ -300,7 +300,7 @@ def get(self, command_method: BoundCommandFunc) -> Cmd2ArgumentParser | None:
 
         return self._parsers.get(full_method_name)
 
-    def remove(self, command_method: BoundCommandFunc) -> None:
+    def remove(self, command_method: BoundCommandFunc[...]) -> None:
         """Remove a given method's parser if it exists."""
         full_method_name = self._fully_qualified_name(command_method)
         if full_method_name in self._parsers:
@@ -919,7 +919,7 @@ def register_command_set(self, cmdset: CommandSet[Any]) -> None:
 
         cmdset.on_register(self)
         methods = cast(
-            list[tuple[str, BoundCommandFunc]],
+            list[tuple[str, BoundCommandFunc[...]]],
             inspect.getmembers(
                 cmdset,
                 predicate=lambda meth: (  # type: ignore[arg-type]
@@ -1021,7 +1021,9 @@ def _build_parser(
 
         return parser
 
-    def _install_command_function(self, command_func_name: str, command_method: BoundCommandFunc, context: str = "") -> None:
+    def _install_command_function(
+        self, command_func_name: str, command_method: BoundCommandFunc[...], context: str = ""
+    ) -> None:
         """Install a new command function into the CLI.
 
         :param command_func_name: name of command function to add
@@ -1086,8 +1088,8 @@ def unregister_command_set(self, cmdset: CommandSet[Any]) -> None:
             cmdset.on_unregister()
             self._unregister_subcommands(cmdset)
 
-            methods: list[tuple[str, BoundCommandFunc]] = cast(
-                list[tuple[str, BoundCommandFunc]],
+            methods: list[tuple[str, BoundCommandFunc[...]]] = cast(
+                list[tuple[str, BoundCommandFunc[...]]],
                 inspect.getmembers(
                     cmdset,
                     predicate=lambda meth: (  # type: ignore[arg-type]
@@ -1161,8 +1163,8 @@ def check_parser_uninstallable(parser: Cmd2ArgumentParser) -> None:
                     )
                 check_parser_uninstallable(subparser)
 
-        methods: list[tuple[str, BoundCommandFunc]] = cast(
-            list[tuple[str, BoundCommandFunc]],
+        methods: list[tuple[str, BoundCommandFunc[...]]] = cast(
+            list[tuple[str, BoundCommandFunc[...]]],
             inspect.getmembers(
                 cmdset,
                 predicate=lambda meth: (  # type: ignore[arg-type]
@@ -3413,17 +3415,17 @@ def _restore_output(self, statement: Statement, saved_redir_state: utils.Redirec
         self._cur_pipe_proc_reader = saved_redir_state.saved_pipe_proc_reader
         self._redirecting = saved_redir_state.saved_redirecting
 
-    def get_command_func(self, command: str) -> BoundCommandFunc | None:
+    def get_command_func(self, command: str) -> BoundCommandFunc[...] | None:
         """Get the bound command function for a command.
 
         :param command: the name of the command
         :return: the bound function implementing the command, or None if not found
         """
         command_func_name = constants.COMMAND_FUNC_PREFIX + command
         command_func = getattr(self, command_func_name, None)
-        return cast(BoundCommandFunc, command_func) if callable(command_func) else None
+        return cast(BoundCommandFunc[...], command_func) if callable(command_func) else None
 
-    def _get_command_category(self, func: BoundCommandFunc) -> str:
+    def _get_command_category(self, func: BoundCommandFunc[...]) -> str:
         """Determine the category for a command.
 
         :param func: the do_* function implementing the command

cmd2/types.py

@@ -69,13 +69,13 @@
 
 # A bound cmd2 command function (e.g. do_command).
 # The 'self' argument is already tied to an instance and is omitted.
-class BoundCommandFunc(Protocol):
+class BoundCommandFunc(Protocol[P]):
     """Protocol for a command function bound to a command instance."""
 
     __name__: str
     __qualname__: str
 
-    def __call__(self, *args: Any, **kwargs: Any) -> bool | None:
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> bool | None:
         """Invoke the bound command function."""
 
 
@@ -94,7 +94,7 @@ def __call__(self, __self: CmdOrSetT, /, *args: P.args, **kwargs: P.kwargs) -> b
     def __get__(self, instance: None, owner: Any = ...) -> "UnboundCommandFunc[CmdOrSetT, P]": ...
 
     @overload
-    def __get__(self, instance: CmdOrSetT, owner: Any = ...) -> BoundCommandFunc: ...
+    def __get__(self, instance: CmdOrSetT, owner: Any = ...) -> BoundCommandFunc[P]: ...
 
 
 ##################################################################################################