Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Express relation between enum and its member in Python typing

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

  • Why MyPy does not state an error for the second call; and
  • How to type 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
like image 613
David Avatar asked Jul 03 '26 18:07

David


1 Answers

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]
like image 178
David Avatar answered Jul 06 '26 06:07

David



Donate For Us

If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!