How to type (in Python, e.g. for MyPy) a function that expects two parameters - an enum and one of its values/members?
from enum import Enum
from typing import TypeVar, Type
class MyEnumA(Enum):
A = 1
B = 2
class MyEnumB(Enum):
A = 1
B = 2
TE = TypeVar('TE', bound=Enum)
def myfunction(member: TE, e: Type[TE]) -> None:
pass
myfunction(MyEnumA.A, MyEnumA) # all right
myfunction(MyEnumA.A, MyEnumB) # I expect mypy-error here but it passed
print(type(MyEnumA.A)) # says: <enum 'MyEnumA'>
print(type(MyEnumB.A)) # says: <enum 'MyEnumB'>
print(f'{isinstance(MyEnumA.A, MyEnumA)=}') # says: isinstance(MyEnumA.A, MyEnumA)=True
print(f'{isinstance(MyEnumA.A, MyEnumB)=}') # says: isinstance(MyEnumA.A, MyEnumB)=False
reveal_type(MyEnumA) # mypy: Revealed type is "def (value: builtins.object) -> e.MyEnumA"
reveal_type(MyEnumA.A) # mypy: Revealed type is "Literal[e.MyEnumA.A]?"
I would like to understand
myfunction
so that MyPy detects an error there.More examples:
myfunction(MyEnumA.A, MyEnumA) # should pass - member of enum
myfunction(MyEnumB.A, MyEnumB) # should pass - member of enum
myfunction(MyEnumA.A, MyEnumB) # should fail - member of other enum
myfunction(MyEnumB.A, MyEnumA) # should fail - member of other enum
The key observation is what MyPy actually considers a class MyEnumA to be:
reveal_type(MyEnumA) # mypy: Revealed type is "def (value: builtins.object) -> e.MyEnumA"
It gives a clue to annotate myfunction
in the following way, no matter how unobvious it is:
P = ParamSpec('P')
def myfunction(member: TE, e: Callable[P, TE]) -> None:
That makes MyPy to find typing error in the second call of myfunction
:
error: Argument 2 to "myfunction" has incompatible type "Type[MyEnumA]"; expected "Callable[[object], MyEnumB]" [arg-type]