How do I keep track of Python enum values internally?

Question

My end goal really is to create a helper method in my Enum class that always returns an Enum member and never raises an exception, given whatever possible value, e.g.

Color.from_value('red')

In case the value is not part of the enum, the helper method will return a default one, say Color.UNKNOWN.

Based on this answer from another related question, I can do checks on the values by listing them through some built-in members. However, what I want to do next is to keep track of all the values in an internal member so that I don't have to always iterate through the values every time the helper method is called. I tried doing something similar to the following:

class Color(Enum):
    RED = 'red'
    BLUE = 'blue'
    GREEN = 'green'

    # this doesn't work
    _values = [item.value for item in Color]

and as expected, it doesn't work. Is this perhaps something that's already built-in in Python Enums?

Answer 1

You can create the method and check for the values in the class:

import enum

class Color(enum.Enum):
    RED = 'red'
    BLUE = 'blue'
    GREEN = 'green'
    UNKNOWN = "unknown"

    @classmethod
    def from_value(cls, value):
      try:
        return {item.value: item for item in cls}[value]
      except KeyError:
        return cls.UNKNOWN

print(Color.from_value("hey"))
print(Color.from_value("red"))

Results:

Color.UNKNOWN
Color.RED

Here you have a live example

In case you dont want to reiterate, you can always have an outside cache of the values:

class Color(enum.Enum):
    RED = 'red'
    BLUE = 'blue'
    GREEN = 'green'
    UNKNOWN = "unknown"

    @classmethod
    def from_value(cls, value):
      try:
        return _COLORITEMS[value]
      except KeyError:
        return cls.UNKNOWN
_COLORITEMS = {item.value: item for item in Color}

Answer 2

You can keep track of values internally, but it's a bit of a hassle:

attempt 1

_values = [k for k in vars().keys() if not k.startswith('_')]

problems 1

>>> # _values is a member
>>> Color._values
<Color._huh: ['RED', 'BLUE', 'GREEN']>

attempt 2

Use Constant from this answer

_values = Constant([k for k in vars().keys() if not k.startswith('_')])

problems 2

It's not really constant as you can still append to the list -- but that can be solved by casting to tuple:

_values = Constant(tuple([k for k in vars().keys() if not k.startswith('_')]))

However, that still doesn't solve the UNKNOWN problem.

Solution

Using either Python 3.6 or aenum 2.0¹ you can specify a _missing_ method that will be called to give your class one last chance before raising ValueError:

class Constant:  # use Constant(object) if in Python 2
    def __init__(self, value):
        self.value = value
    def __get__(self, *args):
        return self.value
    def __repr__(self):
        return '%s(%r)' % (self.__class__.__name__, self.value)


class Color(Enum):

    RED = 'red'
    BLUE = 'blue'
    GREEN = 'green'
    UNKNOWN = 'unknown'
    _values = Constant(tuple([k for k in vars().keys() if not k.startswith('_')]))

    @classmethod
    def _missing_(cls, name):
        if name.lower() in cls._values:
            return cls(name.lower())
        else:
            return cls.UNKNOWN

or

class Color(Enum):

    RED = 'red'
    BLUE = 'blue'
    GREEN = 'green'
    UNKNOWN = 'unknown'

    @classmethod
    def _missing_(cls, name):
        if name == name.lower():
            # avoid infinite recursion
            return cls.UNKNOWN
        else:
            return cls(name.lower())

N.B.: only enum members or None should be returned by _missing_ -- later Python versions will raise TypeError if anything else is returned.

---

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

Answer 3

I have since accepted Netwave's answer but I'm adding this as a bit of variation from the original. I wanted to preserve the answer as is for reference.

import enum

class Color(enum.Enum):
    RED = 'red'
    BLUE = 'blue'
    GREEN = 'green'
    UNKNOWN = "unknown"

    @staticmethod
    def from_value(value):
      try:
        return Color(value)
      except ValueError:
        return cls.UNKNOWN

Notice that in this approach, I have changed @classmethod to @staticmethod as I have no more need for the cls argument. In addition, I'm handling ValueError here instead of KeyError because that's what the Color() raises in case the value couldn't be found. This also works for multi-valued/tuple valued enums.

As for the second approach:

import enum

class Color(enum.Enum):
    RED = 'red'
    BLUE = 'blue'
    GREEN = 'green'
    UNKNOWN = "unknown"

    @staticmethod
    def from_value(value):
        return Color(value) if value in _COLORITEMS else Color.UNKNOWN

_COLORITEMS = [item.value for item in Color]

Here, I switched from dict to list to keep track of values. Instead of handling KeyError, I just simply checked if the value is in the list. I could have used a set but since it's an enum, values are guaranteed to be unique (if I understand correctly).