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!
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.)
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.
For the vast majority of classes (not with Enum
s, 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.
instance_of_a
is A
.A
is not A
— A
is a class, not an instance of a class.type
or instances of custom metaclasses that inherit from type
.A
is type
.The rule is as follows:
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.
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.
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With