Convert enum to Literal type alias in Python typing

user_9541192 0
Viewed 63

Is there a way to type annotate a function or variable in Python in such a way it allows both an enum or Literal formed form the attributes of the enum?

from enum import Enum
from typing import Literal


class State(str, Enum):
    ENABLED = "enabled"
    DISABLED = "disabled"

def is_enabled(state: State | Literal["enabled", "disabled"]) -> bool:
    if isinstance(state, str):
        state = State(state)
    return state == State.ENABLED

In other words, is there a way to obtain the alias for Literal["enabled", "disabled"] without having to rewrite all the keys of the enum?

typing enums python
1 Answers

I'm afraid there's no such way. The first thing that comes to mind is iterating over enum's values to build a Literal type won't work, because Literal cannot contain arbitrary expressions. So, you cannot specify it explicitly:

def is_enabled(state: State | Literal[State.ENABLED.value, State.DISABLED.value]) -> bool:
    ...

There's an open issue on GitHub with the related discussion. Basically, you could have hardcoded literals for every Enum, but in that case, you need to update it in accordance with Enum updates, and one day it will be messy. So, I would either stick with State | str annotation or just State and expect your function to accept only enums.

Also, take into account that you do not need to explicitly create an Enum object to test its value, you can just write "enabled" == State.ENABLED as I mentioned in the comments.

user_99267210
Related