A Hacky way of getting a custom Enum to report the correct type for it value property

[rayansostenes]

What would be a better way to type this without relying on the pyright inference? Is that even possible?

BTW reveal_type(..., expected_text="...") is awesome, but I only discover that it existed after looking at some test files on the repo, I don't think is documented anywhere.

import enum
from typing import TYPE_CHECKING


class ChoiceEnum(enum.Enum):
    def __init__(self, value: str, label: str):
        self._value_ = value
        self._label_: str = label

    @property
    def label(self) -> str:
        return self._label_


class MyChoiceEnum(ChoiceEnum):
    A = ("A", "My A choice")
    B = ("B", "My B choice")
    C = ("C", "My C choice")


assert MyChoiceEnum.A.value == "A"
assert MyChoiceEnum.A.label == "My A choice"

if TYPE_CHECKING:
    reveal_type(MyChoiceEnum.A.value, expected_text="Literal['A']")
    reveal_type(MyChoiceEnum.A.value, expected_text="Literal['B']")
    reveal_type(MyChoiceEnum.A.value, expected_text="Literal['C']")

# My hacky solution to get the value to be a string literal

if TYPE_CHECKING:

    # In this case the inferred type that pyright gives is more accurate
    # than what we could do with a type annotation

    def _(value, label):  # type: ignore
        return value  # type: ignore

else:
    # At runtime we want to return a tuple so that the enum can be instantiated correctly
    def _(value, label):
        return value, label


class MySecondEnum(ChoiceEnum):
    A = _("A", "My A choice")
    B = _("B", "My B choice")
    C = _("C", "My C choice")


assert MySecondEnum.A.value == "A"
assert MySecondEnum.A.label == "My A choice"

if TYPE_CHECKING:
    reveal_type(MySecondEnum.A.value, expected_text="Literal['A']")
    reveal_type(MySecondEnum.B.value, expected_text="Literal['B']")
    reveal_type(MySecondEnum.C.value, expected_text="Literal['C']")

$ npx pyright@latest test.py
No configuration file found.
Assuming Python version 3.9
Assuming Python platform Darwin
Auto-excluding **/node_modules
Auto-excluding **/__pycache__
Auto-excluding **/.*
Searching for source files
Found 1 source file
pyright 1.1.301
test.py
  test.py:25:53 - error: Type text mismatch; expected "Literal['A']" but received "tuple[Literal['A'], Literal['My A choice']]"
  test.py:25:17 - information: Type of "MyChoiceEnum.A.value" is "tuple[Literal['A'], Literal['My A choice']]"
  test.py:26:53 - error: Type text mismatch; expected "Literal['B']" but received "tuple[Literal['A'], Literal['My A choice']]"
  test.py:26:17 - information: Type of "MyChoiceEnum.A.value" is "tuple[Literal['A'], Literal['My A choice']]"
  test.py:27:53 - error: Type text mismatch; expected "Literal['C']" but received "tuple[Literal['A'], Literal['My A choice']]"
  test.py:27:17 - information: Type of "MyChoiceEnum.A.value" is "tuple[Literal['A'], Literal['My A choice']]"
  test.py:52:17 - information: Type of "MySecondEnum.A.value" is "Literal['A']"
  test.py:53:17 - information: Type of "MySecondEnum.B.value" is "Literal['B']"
  test.py:54:17 - information: Type of "MySecondEnum.C.value" is "Literal['C']"
3 errors, 0 warnings, 6 informations
Completed in 0.458sec

[erictraut]

The Enum class in Python is a bit of a nightmare for static type checking because of its nonstandard behaviors and extensibility mechanisms.

My recommendation is to avoid using the more esoteric features of Enum such as the custom values used in the code above. If you decide to use such features, you're going to need to find hacky ways to like this if you want static type checking to work correctly. In most other languages, enums are limited to basic enumerations — for good reason.

Here's a simpler solution that is much more straightforward.

class ChoiceEnum(enum.Enum):
    A = "A"
    B = "B"
    C = "C"

    @property
    def label(self):
        labelMap = {
            "A": "My A choice",
            "B": "My B choice",
            "C": "My C choice",
        }

        return labelMap[self.value]

print(ChoiceEnum.A.label)

[rayansostenes]

Agreed, I have something similar to this on one of my projects right now.

@dataclass
class EnumData:
    label: str
    ...

class MyEnum(enum.Enum):
    A = 'A'
    B = 'B'

    @property
    def data(self) -> EnumData:
        return MY_ENUM_DATA_MAP[self] # pyright: ignore [reportUndefinedVariable]


ENUM_DATA_MAP = {
    MyEnum.A: EnumData(
        label='My Label',
        ...
    ),
    ...
}

I was trying to see if it was possible to improve the django.db.models.Choices types, to actually return the Literal values.