Python Enum with multiple attributes

Question

I have his Enum :

class MultiAttr(Enum):
    RED = (1, 10)
    GREEN = (2, 20)
    BLUE = (3, 20)
    def __init__(self, id, val):
        self.id = id
        self.val = val
assert MultiAttr((2, 20)) == MultiAttr.GREEN

Assertion passes.

1. Why Pylance in VSCode complains Argument missing for parameter "val" ?
2. As the first attribute is the identifer, is there a way to achieve : MultiAttr(2) == MultiAttr.GREEN ?

Answer 1

You could make use of __new__, as you want to customize the actual value of the Enum member:

from enum import Enum

class MultiAttr(bytes, Enum):
    def __new__(cls, value, otherVal):
        obj = bytes.__new__(cls, [value])
        obj._value_ = value
        obj.otherVal = otherVal
        return obj
    RED = (1, 10)
    GREEN = (2, 20)
    BLUE = (3, 20)


print(MultiAttr(2) == MultiAttr.GREEN)
print(MultiAttr(2).otherVal)

Out:

True
20

Answer 2

There is another way to store additional values for enumeration members. extended-enum - helps to do this

from dataclasses import dataclass
from extended_enum import BaseExtendedEnumValue, ExtendedEnum, EnumField

@dataclass(frozen=True)
class SomeExtendedEnumValue(BaseExtendedEnumValue):
    other_value: int

class MultiAttr(ExtendedEnum):
    RED = EnumField(SomeExtendedEnumValue(value=1, other_value=10))
    GREEN = EnumField(SomeExtendedEnumValue(value=2, other_value=20))
    BLUE = EnumField(SomeExtendedEnumValue(value=3, other_value=20))

>>> MultiAttr(SomeExtendedEnumValue(2, 20)) == MultiAttr.GREEN
True
>>> MultiAttr.GREEN.value
2
>>> MultiAttr.GREEN.extended_value
SomeExtendedEnumValue(value=2, other_value=20)
>>> MultiAttr.GREEN.extended_value.other_value
20
>>> MultiAttr(SomeExtendedEnumValue(2, 30)) == MultiAttr.GREEN
ValueError: SomeExtendedEnumValue(value=2, other_value=30) is not a valid MultiAttr

If you don't need to compare the other_value field, you can use field(compare=False). In this case, the comparison will take place on the first field.

from dataclasses import dataclass, field

@dataclass(frozen=True)
class SomeExtendedEnumValue(BaseExtendedEnumValue):
    other_value: int = field(compare=False)

>>> MultiAttr(SomeExtendedEnumValue(2, 30)) == MultiAttr.GREEN
True