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

PLUS: Adding more stuff (a dirt hack) to @Artyer's answer. 🤗

Note that you can also provide "additional" capabilities to an Enum if you create it from a dict, see...

from enum import Enum

_colors = {"RED": (1, "It's the color of blood."), "BLUE": (2, "It's the color of the sky.")}


def _set_members_colors(locals: dict):
    for k, v in colors.items():
        locals[k] = v[0]


class Colors(int, Enum):

    _set_members_colors(locals())

    @property
    def description(self):
        return colors[self.name][1]

print(str(Colors.RED))
print(str(Colors.RED.value))
print(str(Colors.RED.description))

Output...

Colors.RED
1
It's the color of blood.

Thanks! 😉

Answer 2

You can create a new meta class (Either using a meta-metaclass or a factory function, like I do below) that derives from enum.EnumMeta (The metaclass for enums) and just adds the members before creating the class

import enum
import collections.abc


def enum_metaclass_with_default(default_members):
    """Creates an Enum metaclass where `default_members` are added"""
    if not isinstance(default_members, collections.abc.Mapping):
        default_members = enum.Enum('', default_members).__members__

    default_members = dict(default_members)

    class EnumMetaWithDefaults(enum.EnumMeta):
        def __new__(mcs, name, bases, classdict):
            """Updates classdict adding the default members and
            creates a new Enum class with these members
            """

            # Update the classdict with default_members
            # if they don't already exist
            for k, v in default_members.items():
                if k not in classdict:
                    classdict[k] = v

            # Add `enum.Enum` as a base class

            # Can't use `enum.Enum` in `bases`, because
            # that uses `==` instead of `is`
            bases = tuple(bases)
            for base in bases:
                if base is enum.Enum:
                    break
            else:
                bases = (enum.Enum,) + bases

            return super(EnumMetaWithDefaults, mcs).__new__(mcs, name, bases, classdict)

    return EnumMetaWithDefaults


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


class MyEnum(metaclass=enum_metaclass_with_default(value_map)):
    @property
    def double_value(self):
        return self.value * 2


assert MyEnum.a.double_value == 2

---

A different solution was to directly try and update locals(), as it is replaced with a mapping that creates enum values when you try to assign values.

import enum


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


def set_enum_values(locals, value_map):
    # Note that we can't use `locals.update(value_map)`
    # because it's `locals.__setitem__(k, v)` that
    # creates the enum value, and `update` doesn't
    # call `__setitem__`.
    for k, v in value_map:
        locals[k] = v


class MyEnum(enum.Enum):
    set_enum_values(locals(), value_map)

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


assert MyEnum.a.double_value == 2

This seems well defined enough, and a = 1 is most likely going to be the same as locals()['a'] = 1, but it might change in the future. The first solution is more robust and less hacky (And I haven't tested it in other Python implementations, but it probably works the same)

Answer 3

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().