Как расширить возможности стандартного 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