How to define a mapping of enum members in the enum type?

Question

I have a Python enum like this:

from enum import Enum, unique

@unique
class Levels(Enum):
  Unknown = 0
  Warning = 1
  Error =   2

  def __eq__(self, other):    return self.value ==  other.value
  def __ne__(self, other):    return self.value !=  other.value
  def __lt__(self, other):    return self.value <   other.value
  def __le__(self, other):    return self.value <=  other.value
  def __gt__(self, other):    return self.value >   other.value
  def __ge__(self, other):    return self.value >=  other.value

  __MAPPING__ = {
    "warn":   Warning,
    "error":  Error
  }

  @classmethod
  def Parse(cls, value, default=None):
    return cls.__MAPPING__.get(value, default)

In this example, I already extracted the mapping into a class member (according to timeit.timeit(), its faster).

My problem is now, that the dict values are using the original enum member values (integers) instead of the created enum members (EnumMeta). This is reasonable, because the don't exists, when the dict is constructed.

How/Where can I hook into the Enum class / EnumMeta class or my own Levels class, to patch the dict, with the created enum members?

Answer 1

[Only showing the relevant pieces...]

@unique
class Levels(Enum):
    Unknown = 0
    Warning = 1
    Error =   2

    __MAPPING__ = {
        "warn":   Warning,
        "error":  Error
   }

    def __init__(self, *args):
        """Patch the embedded MAP dictionary"""
        self.__class__.__MAPPING__[self._name_] = self

__init__ is called for each member after it has been created, so at that point you can update the __MAPPING__ using the member name as the key. If you use this pattern a lot you should consider using another decorator:

def add_mapping(enum_cls):
    for name in enum_cls.__MAPPING__:
        member = enum_cls.__members__[name]
        enum_cls.__MAPPING__[name] = member
    return enum_cls

and in use:

@add_mapping
@unique
class Levels(Enum):
    Unknown = 0
    Warning = ...

Answer 2

To build upon Mr. Furman's offering above, a map_value(name) method to get the mapping:

add_mapping(enum_cls): 
    enum_cls.__MAPPING__ = {e.value: e for e in enum_cls}
    enum_cls.map_value = lambda name, default=None: enum_cls.__MAPPING__.get(name, default)
    return enum_cls