cliutil: EnumArgument matches user arguments against aliases.

This commit is contained in:
Brett Smith 2021-02-23 14:10:17 -05:00
parent ee2bd6c096
commit 5231a1784f
2 changed files with 13 additions and 7 deletions

View file

@ -94,7 +94,11 @@ class EnumArgument(Generic[ET]):
def enum_type(self, arg: str) -> ET: def enum_type(self, arg: str) -> ET:
"""Return a single enum whose name matches the user argument""" """Return a single enum whose name matches the user argument"""
regexp = re.compile(re.escape(arg), re.IGNORECASE) regexp = re.compile(re.escape(arg), re.IGNORECASE)
matches = frozenset(choice for choice in self.base if regexp.match(choice.name)) matches = frozenset(
choice
for name, choice in self.base.__members__.items()
if regexp.match(name)
)
count = len(matches) count = len(matches)
if count == 1: if count == 1:
return next(iter(matches)) return next(iter(matches))

View file

@ -32,6 +32,8 @@ class ArgChoices(enum.Enum):
AA = 'aa' AA = 'aa'
AB = 'ab' AB = 'ab'
BB = 'bb' BB = 'bb'
START = AA
END = BB
class MockTraceback: class MockTraceback:
@ -270,9 +272,9 @@ def test_diff_year(date, diff, expected):
def test_can_run(cmd, expected): def test_can_run(cmd, expected):
assert cliutil.can_run(cmd) == expected assert cliutil.can_run(cmd) == expected
@pytest.mark.parametrize('choice', ArgChoices) @pytest.mark.parametrize('name,choice', ArgChoices.__members__.items())
def test_enum_arg_enum_type(arg_enum, choice): def test_enum_arg_enum_type(arg_enum, name, choice):
assert arg_enum.enum_type(choice.name) is choice assert arg_enum.enum_type(name.lower()) is choice
assert arg_enum.enum_type(choice.value) is choice assert arg_enum.enum_type(choice.value) is choice
@pytest.mark.parametrize('arg', 'az\0') @pytest.mark.parametrize('arg', 'az\0')
@ -280,9 +282,9 @@ def test_enum_arg_no_enum_match(arg_enum, arg):
with pytest.raises(ValueError): with pytest.raises(ValueError):
arg_enum.enum_type(arg) arg_enum.enum_type(arg)
@pytest.mark.parametrize('choice', ArgChoices) @pytest.mark.parametrize('name,choice', ArgChoices.__members__.items())
def test_enum_arg_value_type(arg_enum, choice): def test_enum_arg_value_type(arg_enum, name, choice):
assert arg_enum.value_type(choice.name) == choice.value assert arg_enum.value_type(name.lower()) == choice.value
assert arg_enum.value_type(choice.value) == choice.value assert arg_enum.value_type(choice.value) == choice.value
@pytest.mark.parametrize('arg', 'az\0') @pytest.mark.parametrize('arg', 'az\0')