How to extend the capabilities of a standard Enum

Or maybe there is still a way to create such an Enum using the standard Python library?!

Sometimes you really want constants to have additional parameters that store other characteristics. The first thing that comes to mind is to describe Enumstoring simple values, and mapping.

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),
}

In this case, it turns out that the constants and characteristics are located on their own, moreover, they can be located in different parts of the system. This can lead to the fact that when a new constant appears in Color, no one will update the mapping, because there is no hard and clear connection.

Let's figure out how you can store everything you need in a single structure.

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

It seems that it looks good, but in this option there is a limitation on the number of additional parameters. Let's try to improve it a little more.

Option 2.

Once in the previous version we managed to do it using tuplethen it will work out with typing.NamedTuple. In addition, there will be named parameters, which will increase readability.

We will store the entire object as a member of the enumeration typing.NamedTuple. Now, in order for us to have a correct comparison of objects, we need to redefine the methods __hash__ And __eq__. Objects will be compared based on one field – 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')
<Color.PURPLE: ColorInfo(value="purple", rgb=RGB(red=128, green=0, blue=128), ru='пурпурный')>

The result was, in principle, a working option. Of course it has its limitations due to the use typing.NamedTuple. Plus, the solution is not universal. After some thought, the following option appeared.

Option 3.

What if instead typing.NamedTuple use dataclass? It seems like a good idea. It becomes possible to inherit classes that store additional. options. Plus helper functions from dataclasses.

As a member of the enumeration, as before, we will store the entire object, only now it is 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__}')

Now for everything to work beautifully, we need a function EnumFieldwhich will simplify initialization (inspired by Pydantic).

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

Now you can start advertising and checking the work.

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='пурпурный')>

Thus the Python package was born extended-enum. In the repository I described main featuresand migration process from standard Enum on ExtendedEnum.

I hope the material was useful! Good luck to you in any endeavor!

PS I will be very pleased if you put a star on github