Search code examples
pythonenumsmypypython-typing

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

Solution

  • 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]