Support argument groups in add_argument_to typing

shtab/__init__.py

@@ -2,8 +2,8 @@
 import logging
 import re
 from argparse import (ONE_OR_MORE, REMAINDER, SUPPRESS, ZERO_OR_MORE, Action, ArgumentParser,
-                      _AppendAction, _AppendConstAction, _CountAction, _HelpAction,
-                      _StoreConstAction, _VersionAction)
+                      _ActionsContainer, _AppendAction, _AppendConstAction, _CountAction,
+                      _HelpAction, _StoreConstAction, _VersionAction)
 from collections import defaultdict
 from contextlib import contextmanager
 from functools import total_ordering
@@ -14,7 +14,7 @@
 from textwrap import dedent
 from typing import Any, Callable
 from typing import Optional as Opt
-from typing import Union
+from typing import TypeVar, Union
 
 try:
     __version__ = version('shtab')
@@ -26,6 +26,7 @@
 
 ShellType = str
 CompleteType = dict[ShellType, Union[str, dict[ShellType, str]]]
+_ActionsContainerT = TypeVar("_ActionsContainerT", bound=_ActionsContainer)
 SUPPORTED_SHELLS: list[ShellType] = []
 _SUPPORTED_COMPLETERS: dict[ShellType, Callable] = {}
 CHOICE_FUNCTIONS: dict[str, CompleteType] = {
@@ -1508,12 +1509,12 @@ def __call__(self, parser, namespace, values, option_string=None):
 
 
 def add_argument_to(
-    parser: ArgumentParser,
+    parser: _ActionsContainerT,
     option_string: Union[str, list[str]] = "--print-completion",
     help: str = "print shell completion script",                    # pylint: disable=W0622
     parent: Opt[ArgumentParser] = None,
     preamble: Union[str, dict[str, str]] = "",
-):
+) -> _ActionsContainerT:
     """
     option_string:
       iff positional (no `-` prefix) then `parser` is assumed to actually be

tests/test_shtab.py

@@ -760,6 +760,13 @@ def test_add_argument_to_optional(shell):
         pytest.skip("WiP")
 
 
+def test_add_argument_to_group_typing() -> None:
+    """Ensure mypy accepts argument groups and the function preserves identity."""
+    parser = ArgumentParser(prog="test")
+    group = parser.add_argument_group("completion")
+    assert shtab.add_argument_to(group) is group
+
+
 @fix_shell
 def test_add_argument_to_positional(shell, capsys):
     parser = ArgumentParser(prog="test")