Skip to content
Advertisement

Type hint for return value in subclass

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.

User contributions licensed under: CC BY-SA
6 People found this is helpful
Advertisement