Как расширить возможности стандартного 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('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('black')
Получился в принципе рабочий вариант. Конечно у него есть свои ограничения за счет использования 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.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')
Так родился Python пакет extended-enum. В репозитории я описал основные возможности, а также процесс миграции со стандартного Enum на ExtendedEnum.
Надеюсь, что материал был полезен! Успехов вам в любых начинаниях!
P.S. Мне будет очень приятно, если поставите звёздочку на github
