Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

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!

like image 724
HitLuca Avatar asked Jul 07 '21 08:07

HitLuca


2 Answers

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 AA 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.

like image 65
Alex Waygood Avatar answered Oct 17 '22 03:10

Alex Waygood


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.

like image 1
Clément Avatar answered Oct 17 '22 01:10

Clément