Can one add a custom error in Enum to show valid values?

Viewed 353

Let's say I have

from enum import Enum

class SomeType(Enum):
    TYPEA = 'type_a'
    TYPEB = 'type_b'
    TYPEC = 'type_c'

If I now do

SomeType('type_a')

I will get

<SomeType.TYPEA: 'type_a'>

as expected. When I do

SomeType('type_o')

I will receive

ValueError: 'type_o' is not a valid SomeType

which is also expected.

My question is: Can one somehow easily customize the error so that it shows all valid types? So, in my case I would like to have

ValueError: 'type_o' is not a valid SomeType. Valid types are 'type_a', 'type_b', 'type_c'.

1 Answers

Use the _missing_ method:

from enum import Enum

class SomeType(Enum):
    TYPEA = 'type_a'
    TYPEB = 'type_b'
    TYPEC = 'type_c'
    @classmethod
    def _missing_(cls, value):
        raise ValueError(
                '%r is not a valid %s.  Valid types: %s' % (
                    value,
                    cls.__name__,
                    ', '.join([repr(m.value) for m in cls]),
                    ))

and in use:

>>> SomeType('type_a')
<SomeType.TYPEA: 'type_a'>

>>> SomeType('type_o')
ValueError: 'type_o' is not a valid SomeType

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  ...
ValueError: 'type_o' is not a valid SomeType.  Valid types: 'type_a', 'type_b', 'type_c'

As you see, it's a little clunky with the exception chaining as Enum itself will raise the "primary" ValueError, with your missing error being in the chain. Depending on your needs, you can narrow the ValueError that you are raising in _missing_ to just include the valid types:

from enum import Enum

class SomeType(Enum):
    TYPEA = 'type_a'
    TYPEB = 'type_b'
    TYPEC = 'type_c'
    @classmethod
    def _missing_(cls, value):
        raise ValueError(
                'Valid types: %s' % (
                    ', '.join([repr(m.value) for m in cls]),
                    )

and in use:

>>> SomeType('type_o')
ValueError: 'type_o' is not a valid SomeType

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  ...
ValueError: Valid types: 'type_a', 'type_b', 'type_c'

Disclosure: I am the author of the Python stdlib Enum, the enum34 backport, and the Advanced Enumeration (aenum) library.

Related