From 5231a1784f8795669fdba3b7eaf5881326ae89b4 Mon Sep 17 00:00:00 2001 From: Brett Smith Date: Tue, 23 Feb 2021 14:10:17 -0500 Subject: [PATCH] cliutil: EnumArgument matches user arguments against aliases. --- conservancy_beancount/cliutil.py | 6 +++++- tests/test_cliutil.py | 14 ++++++++------ 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/conservancy_beancount/cliutil.py b/conservancy_beancount/cliutil.py index 7ea8147..3e7138e 100644 --- a/conservancy_beancount/cliutil.py +++ b/conservancy_beancount/cliutil.py @@ -94,7 +94,11 @@ class EnumArgument(Generic[ET]): def enum_type(self, arg: str) -> ET: """Return a single enum whose name matches the user argument""" 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) if count == 1: return next(iter(matches)) diff --git a/tests/test_cliutil.py b/tests/test_cliutil.py index b4013b6..510d615 100644 --- a/tests/test_cliutil.py +++ b/tests/test_cliutil.py @@ -32,6 +32,8 @@ class ArgChoices(enum.Enum): AA = 'aa' AB = 'ab' BB = 'bb' + START = AA + END = BB class MockTraceback: @@ -270,9 +272,9 @@ def test_diff_year(date, diff, expected): def test_can_run(cmd, expected): assert cliutil.can_run(cmd) == expected -@pytest.mark.parametrize('choice', ArgChoices) -def test_enum_arg_enum_type(arg_enum, choice): - assert arg_enum.enum_type(choice.name) is choice +@pytest.mark.parametrize('name,choice', ArgChoices.__members__.items()) +def test_enum_arg_enum_type(arg_enum, name, choice): + assert arg_enum.enum_type(name.lower()) is choice assert arg_enum.enum_type(choice.value) is choice @pytest.mark.parametrize('arg', 'az\0') @@ -280,9 +282,9 @@ def test_enum_arg_no_enum_match(arg_enum, arg): with pytest.raises(ValueError): arg_enum.enum_type(arg) -@pytest.mark.parametrize('choice', ArgChoices) -def test_enum_arg_value_type(arg_enum, choice): - assert arg_enum.value_type(choice.name) == choice.value +@pytest.mark.parametrize('name,choice', ArgChoices.__members__.items()) +def test_enum_arg_value_type(arg_enum, name, choice): + assert arg_enum.value_type(name.lower()) == choice.value assert arg_enum.value_type(choice.value) == choice.value @pytest.mark.parametrize('arg', 'az\0')