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:
Self = TypeVar("Self", bound="Shape") class Shape: @classmethod def from_config(cls: type[Self], config: dict[str, float]) -> Self: return cls(config["scale"])We propose using
Self
directly:from typing import Self class Shape: @classmethod def from_config(cls, config: dict[str, float]) -> Self: return cls(config["scale"])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.