Created
March 31, 2023 23:32
-
-
Save rayansostenes/a40e9d890512a436af031fb0d9616bc4 to your computer and use it in GitHub Desktop.
A Hacky way of getting a custom Enum to report the correct type for it value property
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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']") | |
OUTPUT = """ | |
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 | |
""" |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment