I am writing a CustomEnum class in which I want to add some helper methods, that would then be available by the classes subclassing my CustomEnum. One of the methods is to return a random enum value, and this is where I am stuck. The function works as expected, but on the type-hinting side, I cannot figure out a way of saying “the return type is the same type of cls”.
I am fairly sure there’s some TypeVar
or similar magic involved, but since I never had to use them I never took the time to figure them out.
class CustomEnum(Enum):
@classmethod
def random(cls) -> ???:
return random.choice(list(cls))
class SubclassingEnum(CustomEnum):
A = "a"
B = "b"
random_subclassing_enum: SubclassingEnum
random_subclassing_enum = SubclassingEnum.random() # Incompatible types in assignment (expression has type "CustomEnum", variable has type "SubclassingEnum")
Can somebody help me or give me a hint on how to proceed?
Thanks!
Advertisement
Answer
Starting in Python 3.11, the correct return annotation for this code is Self
:
from typing import Self
class CustomEnum(Enum):
@classmethod
def random(cls) -> Self:
return random.choice(list(cls))
Quoting from the PEP:
This PEP introduces a simple and intuitive way to annotate methods that return an instance of their class. This behaves the same as the TypeVar-based approach specified in PEP 484 but is more concise and easier to follow.
The current workaround for this is unintuitive and error-prone:
JavaScript161Self = TypeVar("Self", bound="Shape")
2class Shape:
3@classmethod
4def from_config(cls: type[Self], config: dict[str, float]) -> Self:
5return cls(config["scale"])
6
We propose using
Self
directly:JavaScript161from typing import Self
2class Shape:
3@classmethod
4def from_config(cls, config: dict[str, float]) -> Self:
5return cls(config["scale"])
6
This avoids the complicated
cls: type[Self]
annotation and the TypeVar declaration with a bound. Once again, the latter code behaves equivalently to the former code.