А может всё-таки есть способ сделать такой 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
