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?
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:
# THIS DOES NOT WORK
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.