diff --git a/conservancy_beancount/cliutil.py b/conservancy_beancount/cliutil.py index 7ea814766869ef346b8221b164c4d2a84b9bb78b..3e7138e3fb5eac5f4db1b574de62060e927df9af 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 b4013b6777c1d5c3a8ade47e1c5dc10aed736c19..510d6152c6d8d2f9ca7723a379d580c018774a9f 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')