diff --git a/shtab/__init__.py b/shtab/__init__.py index 2f3ed2d..b4cb8c4 100644 --- a/shtab/__init__.py +++ b/shtab/__init__.py @@ -20,7 +20,7 @@ from string import Template from typing import Any, Dict, List from typing import Optional as Opt -from typing import Sequence, Union +from typing import Union # version detector. Precedence: installed dist, git, 'UNKNOWN' try: @@ -112,8 +112,8 @@ def __lt__(self, other: object) -> bool: class CustomChoices: - def __init__(self, choices: Sequence[str]) -> None: - self.choices = tuple(choices) + def __init__(self, *choices: str) -> None: + self.choices = choices assert all(isinstance(c, str) for c in self.choices) def bash_choices(self, prefix: str, option_string: str) -> str: @@ -128,8 +128,8 @@ def tcsh_choices(self) -> str: return "(" + " ".join(self.choices) + ")" -def custom_choices(choices: List[str]) -> CustomChoices: - return CustomChoices(choices) +def custom_choices(*choices: str) -> CustomChoices: + return CustomChoices(*choices) class Optional: diff --git a/tests/test_shtab.py b/tests/test_shtab.py index 94c99b6..bd671f4 100644 --- a/tests/test_shtab.py +++ b/tests/test_shtab.py @@ -202,7 +202,7 @@ def test_subparser_custom_complete(shell, caplog): def test_custom_choices(shell, caplog): parser = ArgumentParser(prog="test") action = parser.add_argument("--optA") - action.complete = shtab.custom_choices(["yes", "no", "maybe"]) + action.complete = shtab.custom_choices("yes", "no", "maybe") with caplog.at_level(logging.INFO): completion = shtab.complete(parser, shell=shell) print(completion)