python - Type hint for return value in subclass

Question

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!

Answer 1

The syntax here is kind of horrible, but I don't think there's a cleaner way to do this. The following passes MyPy:

from typing import TypeVar
from enum import Enum
import random 

T = TypeVar("T", bound="CustomEnum")

class CustomEnum(Enum):
    @classmethod
    def random(cls: type[T]) -> T:
        return random.choice(list(cls))

(In python versions <= 3.8, you have to use typing.Type rather than the builtin type if you want to parameterise it.)

What's going on here?

T is defined at the top as being a type variable that is "bound" to the CustomEnum class. This means that a variable annotated with T can only be an instance of CustomEnum or an instance of a class inheriting from CustomEnum.

In the classmethod above, we're actually using this type-variable to define the type of the cls parameter with respect to the return type. Usually we do the opposite — we usually define a function's return types with respect to the types of that function's input parameters. So it's understandable if this feels a little mind-bending!

We're saying: this method leads to instances of a class — we don't know what the class will be, but we know it will either be CustomEnum or a class inheriting from CustomEnum. We also know that whatever class is returned, we can guarantee that the type of the cls parameter in the function will be "one level up" in the type heirarchy from the type of the return value.

In a lot of situations, we might know that type[cls] will always be a fixed value. In those situations, it would be possible to hardcode that into the type annotations. However, it's best not to do so, and instead to use this method, which clearly shows the relationship between the type of the input and the return type (even if it uses horrible syntax to do so!).

Further reading: the MyPy documentation on the type of class objects.

Further explanation and examples

For the vast majority of classes (not with Enums, they use metaclasses, but let's leave that aside for the moment), the following will hold true:

Example 1

Class A:
    pass

instance_of_a = A()
type(instance_of_a) == A # True
type(A) == type # True

Example 2

class B:
    pass

instance_of_b = B()
type(instance_of_b) == B # True
type(B) == type # True

For the cls parameter of your CustomEnum.random() method, we're annotating the equivalent of A rather than instance_of_a in my Example 1 above.

• The type of instance_of_a is A.
• But the type of A is not A — A is a class, not an instance of a class.
• Classes are not instances of classes; they are either instances of type or instances of custom metaclasses that inherit from type.
• No metaclasses are being used here; ergo, the type of A is type.

The rule is as follows:

• The type of all python class instances will be the class they're an instance of.
• The type of all python classes will be either type or (if you're being too clever for your own good) a custom metaclass that inherits from type.

With your CustomEnum class, we could annotate the cls parameter with the metaclass that the enum module uses (enum.EnumType, if you want to know). But, as I say — best not to. The solution I've suggested illustrates the relationship between the input type and the return type more clearly.