CODEWITHSUNDEEP
pythonenums
ac24
ac24

Reputation: 5565

Python enum from string, raise error if string not in members

I am wondering if I can use python enums to enforce specific string inputs from a user into my functions and classes, and avoid having to code the explicit checks against acceptable values.

So instead of this:

e.g.

utilities_supplied = {
    'electricity': 'Yes',
    'gas': 'No',
    'water': 'Yes',
}

def find_utility(utility):
    try:
        print(utilities_supplied[utility])
    except KeyError:
        raise KeyError(f'Please choose one of {utilities_supplied.keys()}, \'{utility}\' provided')

I would like to do this:


from enum import Enum
class Utility(Enum):
    electricity = 1
    gas = 2
    water = 3

Utility['electric']  
# not a member, I want this to raise an error which lists the acceptable options.

Is it possible for enums to raise an error listing the possible enumeration members?

Upvotes: 1

Views: 5350

Answers (2)

ac24
ac24

Reputation: 5565

EDIT: This doesn't really work when I subclass the ValidatedEnum class. Any help gratefully received.

The solution using _missing_ was fairly straightforward it turns out! Thanks for the idea @EthanFurman.

from enum import Enum
class ValidatedEnum(Enum):
    electricity = 1
    gas = 2
    water = 3

    @classmethod
    def _missing_(cls, value):
        choices = list(cls.__members__.keys())
        raise ValueError("%r is not a valid %s, please choose from %s" % (value, cls.__name__, choices))


<ipython-input-1-8b49d805ac2d> in _missing_(cls, value)
      8     def _missing_(cls, value):
      9         choices = list(cls.__members__.keys())
---> 10         raise ValueError("%r is not a valid %s, please choose from %s" % (value, cls.__name__, choices))
     11 

ValueError: 'electric' is not a valid ValidatedEnum, please choose from ['electricity', 'gas', 'water']

Upvotes: 1

Maurice Meyer
Maurice Meyer

Reputation: 18106

You can use EnumMeta to overwrite __getitem__, which fetches the values from the enum:

import enum


class MyEnum(enum.EnumMeta):
    def __getitem__(cls, name):
        try:
            return super().__getitem__(name)
        except (KeyError) as error:
            options = ', '.join(cls._member_map_.keys())
            msg = f'Please choose one of {options}, \'{name}\' provided'
            raise ValueError(msg) from None


class Utility(enum.Enum, metaclass=MyEnum):
    electricity = 1
    gas = 2
    water = 3


fire = Utility['fire']
print(fire)

Output:

ValueError: Please choose one of electricity, gas, water, 'fire' provided

Upvotes: 2

Related Questions