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.

Lodgings

Can't you just do Enum('MyEnum', {k: v * 2 for k, v in value_map.items()})? – Egress

I can, but this double_value is just a dummy example, I would like to be able to define any methods. – Lodgings

I'm not sure what is stopping you from defining such methods? please provide an example of what you are actually trying to achieve – Egress

I don't understand, how you can just use self.value on that method. Since there is no instance variable named value :/ – Ball

@Egress What stopps me, is that if I use this: MyEnum = Enum('MyEnum', value_map), then I cannot define methods. But if I use this class MyEnum(Enum): ..., then I must write all values explicitely, and cannot use my pre-defined dict. – Lodgings

@hansolo the MyEnum class inherits from Enum, and then it will have a value and name as well. That is what enum does. – Lodgings

@Egress I think of somehow overriding the __prepare__ or __new__ magic methods of Enum and use my dict there. – Lodgings

@hansolo I guess so, but I don't know how yet :) That is why I'm asking :) – Lodgings

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

Lulululuabourg

That's nice, thanks! Any possibility to do this without having to use this enum_with_extras function? I mean as a class definition somehow? – Lodgings

@waszil: I took a guess at what you might want and added that as an option: adding in a base enum class without members. – Lulululuabourg

This type argument method is what I thought of, thank you very much! – Lodgings

One more question: if I would like to use this schema for IntEnum, I get: TypeError: object.__new__(MyEnum) is not safe, use int.__new__() – Lodgings

Got it, DoubledEnum has to be subclassed from enum.IntEnum as well. – Lodgings

@waszil: that, or make sure you use the correct ordering for base classes. – Lulululuabourg

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)

Gauze

Answer 3

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! 😉

Haro