Как создать класс Python Enum из существующего dict с помощью дополнительных методов?

Question

Допустим, у меня есть уже существующее отображение в качестве словаря:

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

Из этого я могу создать класс enum следующим образом:

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

и пользуйся вот так

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

Но затем я хочу определить некоторые методы для моего нового класса enum:

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

Конечно, я могу сделать это:

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

Но, как я уже сказал, мне нужно использовать предопределенный словарь отображения значений, поэтому я не могу этого сделать. Как этого достичь? Я пытался унаследовать от другого класса, определяющего этот метод как миксин, но я не мог понять это.

Answer 1

Вы можете передать базовый тип с помощью методов mixin в функциональный API с аргументом type:

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

Для полностью функционального подхода, в котором никогда не используется оператор class, вы можете создать базовое сочетание с функцией type() :

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

Вы также можете использовать метакласс enum.EnumMeta() таким же образом, как Python при создании подкласса class MyEnum(enum.Enum): ...:

1. Создание словаря классов с использованием метакласса __prepare__ hook
2. Вызовите метакласс, передав ему имя класса, базы ((enum.Enum,) здесь) и словарь классов, созданный на шаге 1.

Пользовательский подкласс словаря, который используется enum.EnumMeta, не предназначен для простого повторного использования; он реализует хук __setitem__ для записи метаданных, но не переопределяет метод dict.update(), поэтому мы должны быть немного осторожнее при использовании словаря value_map:

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)

Затем передайте double_value=property(double_value) этой функции (вместе с именем enum и словарем value_map):

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

В противном случае вам разрешено создавать подклассы перечисления без членов (все, что является дескриптором 1039 *, не является членом, поэтому функции, свойства, методы класса и т. Д.), Поэтому вы можете сначала определить перечисление без элементов:

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

, который является приемлемым базовым классом как для функционального API (например, enum.Enum(..., type=DoubledEnum)), так и для подхода метакласса, который я кодировал как enum_with_extras().

Answer 2

Вы можете создать новый мета класс (используя мета-метакласс или фабричную функцию, как я это делаю ниже), который наследуется от enum.EnumMeta (метакласс для перечислений) и просто добавляет членов перед созданием класса * 1002.*

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

---

Другое решение состояло в том, чтобы напрямую попытаться обновить locals(), поскольку оно заменяется отображением, которое создает значения перечисления при попытке присвоить значения.

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

Кажется, это достаточно четко определено, и a = 1, скорее всего, будет таким же, как locals()['a'] = 1, но может измениться в будущем.Первое решение более надежное и менее хакерское (и я не тестировал его в других реализациях Python, но, вероятно, оно работает так же)