How do I type hint for Enums in Python?

Question

I have a python function for which I want to use type hinting. There are two arguments. The first is any Enum class, the second optional arg is an element of that Enum.

For example, say I have:

class Foo(Enum):
    ALPHA = 1
    BETA = 2
    GAMMA = 3

The first arg would be, e.g. Foo, the second would be e.g. Foo.ALPHA

What would be the correct way of type hinting this? What I have so far is:

def switch(options: Enum, selected: Optional[Enum] = None) -> Enum:
   # Rest of fn...

but that doesn't seem right.

Answer 1

Define a TypeVar with Enum as a bound, and then specify that your function takes the Type of that typevar and returns an instance of it:

from enum import Enum
from typing import Optional, Type, TypeVar

_E = TypeVar('_E', bound=Enum)

def switch(
    options: Type[_E],
    selected: Optional[_E] = None
) -> _E:
    ...

Testing it in mypy with an actual Enum subclass:

class Foo(Enum):
    ALPHA = 1
    BETA = 2
    GAMMA = 3

reveal_type(switch(Foo, Foo.ALPHA))  # Revealed type is "Foo"