Как расширить возможности стандартного Enum / Хабр

А может всё-таки есть способ сделать такой Enum, используя стандартную библиотеку Python?!

Иногда очень хочется, чтобы у констант были дополнительные параметры, хранящие прочие характеристики. Первое, что приходит на ум - это описать Enum, хранящий простые значения, и маппинг.

from dataclasses import dataclass
from enum import Enum

class Color(Enum):
    BLACK = 'black'
    WHITE = 'white'
    PURPLE = 'purple'
    

@dataclass(frozen=True)
class RGB:
    red: int
    green: int
    blue: int


COLOR_TO_RGB = {
    Color.BLACK: RGB(0, 0, 0),
    Color.WHITE: RGB(255, 255, 255),
    Color.PURPLE: RGB(128, 0, 128),
}

В таком случае получается, что константы и характеристики располагаются сами по себе, к тому же они могут находится в разных частях системы. Это может привести к тому что при появлении новой константы в Color, никто не обновит маппинг, т.к. нет жёсткой и явной связи.

Давайте разбираться как же можно хранить всё необходимое в единой структуре.

Вариант 1.

from enum import Enum
from typing import Union

class RelatedEnum(str, Enum):
    related_value: Union[int, str]

    def __new__(
            cls, 
            value: Union[int, str], 
            related_value: Union[int, str]
    ) -> 'RelatedEnum':
        obj = str.__new__(cls, value)
        obj._value_ = value
        obj.related_value = related_value
        return obj

class SomeEnum(RelatedEnum):
    CONST1 = ('value1', 'related_value1')
    CONST2 = ('value2', 'related_value2')

>>> SomeEnum.CONST1.value
'value1'
>>> SomeEnum.CONST1.related_value
'related_value1'
>>> SomeEnum('value1')
<SomeEnum.CONST1: 'value1'>
>>> SomeEnum('value1').related_value
'related_value1'

Кажется, что выглядит неплохо, но в таком варианте есть ограничение по кол-ву дополнительных параметров. Давайте попробуем еще немного улучшить.

Вариант 2.

Раз в прошлом варианте у нас получилось сделать с использованием tuple, то значит получится и с typing.NamedTuple. К тому же будут именованные параметры, что повысит читабельность.

В качестве члена перечисления будем хранить целиком объект typing.NamedTuple. Теперь чтобы у нас происходило корректное сравнение объектов нам нужно переопределить методы __hash__ и __eq__. Сравниваться объекты будут по одному полю - value.

from enum import Enum
from types import DynamicClassAttribute
from typing import Any, NamedTuple, Union


class RGB(NamedTuple):
    red: int
    green: int
    blue: int


class ColorInfo(NamedTuple):
    value: Union[int, str]
    rgb: RGB = None
    ru: str = None

    def __hash__(self) -> int:
        return hash(self.value)

    def __eq__(self, other: Any) -> bool:
        if isinstance(other, type(self)):
            return hash(self) == hash(other)
        return False


class Color(Enum):
    BLACK = ColorInfo('black', rgb=RGB(0, 0, 0), ru='черный')
    WHITE = ColorInfo('white', rgb=RGB(255, 255, 255), ru='белый')
    RED = ColorInfo('red', rgb=RGB(255, 0, 0), ru='красный')
    GREEN = ColorInfo('green', rgb=RGB(0, 255, 0), ru='зеленый')
    BLUE = ColorInfo('blue', rgb=RGB(0, 0, 255), ru='голубой')
    PURPLE = ColorInfo('purple', rgb=RGB(128, 0, 128), ru='пурпурный')
    OLIVE = ColorInfo('olive', rgb=RGB(128, 128, 0), ru='оливковый')
    TEAL = ColorInfo('teal', rgb=RGB(0, 128, 128), ru='бирюзовый')

    _value_: ColorInfo

    @DynamicClassAttribute
    def value(self) -> str:
        return self._value_.value

    @DynamicClassAttribute
    def info(self) -> ColorInfo:
        return self._value_

    @classmethod
    def _missing_(cls, value: Any) -> 'Color':
        if isinstance(value, (str, int)):
            return cls._value2member_map_[ColorInfo(value)]
        raise ValueError(f'Unknown color: {value}')

>>> Color.PURPLE.value
'purple'
>>> Color.PURPLE.info
ColorInfo(value='purple', rgb=RGB(red=128, green=0, blue=128), ru='пурпурный')
>>> Color('purple')
<Color.PURPLE: ColorInfo(value='purple', rgb=RGB(red=128, green=0, blue=128), ru='пурпурный')>

Получился в принципе рабочий вариант. Конечно у него есть свои ограничения за счет использования typing.NamedTuple. Плюс ко всему решение не универсальное. После некоторых раздумий появился следующий вариант.

Вариант 3.

Что если вместо typing.NamedTuple использовать dataclass? Вроде идея здравая. Появляется возможность наследования классов, хранящих доп. параметры. Плюс вспомогательные функции из dataclasses.

В качестве члена перечисления, как и в прошлый раз, будем хранить объект целиком, только теперь это dataclass.

import enum
from dataclasses import dataclass
from types import DynamicClassAttribute
from typing import Union, TypeVar, Any
from uuid import UUID

SimpleValueType = Union[UUID, int, str]
ExtendedEnumValueType = TypeVar('ExtendedEnumValueType', bound='BaseExtendedEnumValue')
ExtendedEnumType = TypeVar('ExtendedEnumType', bound='ExtendedEnum')

@dataclass(frozen=True)
class BaseExtendedEnumValue:
    value: SimpleValueType
    
class ExtendedEnum(enum.Enum):
    value: SimpleValueType
    _value_: ExtendedEnumValueType

    @DynamicClassAttribute
    def value(self) -> SimpleValueType:
        return self._value_.value

    @DynamicClassAttribute
    def extended_value(self) -> ExtendedEnumValueType:
        return self._value_

    @classmethod
    def _missing_(cls, value: Any) -> ExtendedEnumType:  # noqa: WPS120
        if isinstance(value, (UUID, int, str)):
            simple_value2member = {member.value: member for member in cls.__members__.values()}
            try:
                return simple_value2member[value]
            except KeyError:
                pass  # noqa: WPS420
        raise ValueError(f'{value!r} is not a valid {cls.__qualname__}')

Теперь чтобы всё красиво заработало нам понадобится функция EnumField, которая упростит инициализацию (вдохновлено Pydantic).

def EnumField(value: Union[SimpleValueType, ExtendedEnumValueType]) -> BaseExtendedEnumValue:
    if isinstance(value, (UUID, int, str)):
        return BaseExtendedEnumValue(value=value)
    return value

Теперь можно приступать к объявлению и проверке работы.

from dataclasses import field

@dataclass(frozen=True)
class RGB:
    red: int
    green: int
    blue: int

@dataclass(frozen=True)
class ColorInfo(BaseExtendedEnumValue):
    rgb: RGB = field(compare=False)
    ru: str = field(compare=False)

class Color(ExtendedEnum):
    BLACK = EnumField(ColorInfo('black', rgb=RGB(0, 0, 0), ru='черный'))
    WHITE = EnumField(ColorInfo('white', rgb=RGB(255, 255, 255), ru='белый'))
    RED = EnumField(ColorInfo('red', rgb=RGB(255, 0, 0), ru='красный'))
    GREEN = EnumField(ColorInfo('green', rgb=RGB(0, 255, 0), ru='зеленый'))
    BLUE = EnumField(ColorInfo('blue', rgb=RGB(0, 0, 255), ru='голубой'))
    PURPLE = EnumField(ColorInfo('purple', rgb=RGB(128, 0, 128), ru='пурпурный'))
    OLIVE = EnumField(ColorInfo('olive', rgb=RGB(128, 128, 0), ru='оливковый'))
    TEAL = EnumField(ColorInfo('teal', rgb=RGB(0, 128, 128), ru='бирюзовый'))

>>> Color.PURPLE
<Color.PURPLE: ColorInfo(value='purple', rgb=RGB(red=128, green=0, blue=128), ru='пурпурный')>
>>> Color.PURPLE.value
'purple'
>>> Color.PURPLE.extended_value
ColorInfo(value='purple', rgb=RGB(red=128, green=0, blue=128), ru='пурпурный')
>>> Color.PURPLE.extended_value.rgb
RGB(red=128, green=0, blue=128)
>>> Color.PURPLE.extended_value.ru
'пурпурный'
>>> Color('purple')
<Color.PURPLE: ColorInfo(value='purple', rgb=RGB(red=128, green=0, blue=128), ru='пурпурный')>

Так родился Python пакет extended-enum. В репозитории я описал основные возможности, а также процесс миграции со стандартного Enum на ExtendedEnum.

Надеюсь, что материал был полезен! Успехов вам в любых начинаниях!

P.S. Мне будет очень приятно, если поставите звёздочку на github

Как расширить возможности стандартного Enum

Публикации

Ближайшие события