diff --git a/shtab/__init__.py b/shtab/__init__.py index 641d5c1..77ab438 100644 --- a/shtab/__init__.py +++ b/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 diff --git a/tests/test_shtab.py b/tests/test_shtab.py index 0692b7b..d834d4c 100644 --- a/tests/test_shtab.py +++ b/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")