How to create Python Enum class from existing dict with additional methods?

Question

Let's say, I have a pre-existing mapping as a dictionary:

value_map = {'a': 1, 'b': 2}

I can create an enum class from this like so:

from enum import Enum
MyEnum = Enum('MyEnum', value_map)

and use it like so

a = MyEnum.a
print(a.value)
>>> 1
print(a.name)
>>> 'a'

But then I want to define some methods to my new enum class:

def double_value(self):
    return self.value * 2

Of course, i can do this:

class MyEnum(Enum):
    a = 1
    b = 2
    @property
    def double_value(self):
        return self.value * 2

But as I said, I have to use a pre-defined value mapping dictionary, so I cannot do this. How can this be achieved? I tried to inherit from another class defining this method like a mixin, but I could'nt figure it out.

Answer 1

You can pass in a base type with mixin methods into the functional API, with the type argument:

>>> import enum
>>> value_map = {'a': 1, 'b': 2}
>>> class DoubledEnum:
...     @property
...     def double_value(self):
...         return self.value * 2
...
>>> MyEnum = enum.Enum('MyEnum', value_map, type=DoubledEnum)
>>> MyEnum.a.double_value
2

For a fully functional approach that never uses a class statement, you can create the base mix-in with the type() function:

DoubledEnum = type('DoubledEnum', (), {'double_value': property(double_value)})
MyEnum = enum.Enum('MyEnum', value_map, type=DoubledEnum)

You can also use enum.EnumMeta() metaclass the same way, the way Python would when you create a class MyEnum(enum.Enum): ... subclass:

1. Create a class dictionary using the metaclass __prepare__ hook
2. Call the metaclass, passing in the class name, the bases ((enum.Enum,) here), and the class dictionary created in step 1.

The custom dictionary subclass that enum.EnumMeta uses isn't really designed for easy reuse; it implements a __setitem__ hook to record metadata, but doesn't override the dict.update() method, so we need to use a little care when using your value_map dictionary:

import enum

def enum_with_extras(name, value_map, bases=enum.Enum, **extras):
    if not isinstance(bases, tuple):
        bases = bases,
    if not any(issubclass(b, enum.Enum) for b in bases):
        bases += enum.Enum,
    classdict = enum.EnumMeta.__prepare__(name, bases)
    for key, value in {**value_map, **extras}.items():
        classdict[key] = value
    return enum.EnumMeta(name, bases, classdict)

Then pass in double_value=property(double_value) to that function (together with the enum name and value_map dictionary):

>>> def double_value(self):
...     return self.value * 2
...
>>> MyEnum = enum_with_extras('MyEnum', value_map, double_value=property(double_value))
>>> MyEnum.a
<MyEnum.a: 1>
>>> MyEnum.a.double_value
2

You are otherwise allowed to create subclasses of an enum without members (anything that's a descriptor is not a member, so functions, properties, classmethods, etc.), so you can define an enum without members first:

class DoubledEnum(enum.Enum):
    @property
    def double_value(self):
        return self.value * 2

which is an acceptable base class for both in the functional API (e.g. enum.Enum(..., type=DoubledEnum)) and for the metaclass approach I encoded as enum_with_extras().