When defining a function, there are several ways to limit input to a set of predefined options. I thought it would make sense to use Enum
objects to achieve this. For example:
from enum import Enum, auto
class ColorOptions(Enum):
RED = auto()
BLUE = auto()
def color_something(color: ColorOptions):
match color:
case ColorOptions.RED:
return 'rgb(1, 0, 0)'
case ColorOptions.BLUE:
return 'rgb(0, 0, 1)'
# Example usage
print(color_something(ColorOptions.BLUE)) # Output: 'rgb(0, 0, 1)'
However, the following call doesn’t work:
color_something('BLUE')
To address this, I modified the function to accept both Enum
members and string representations, like so:
def color_something(color: ColorOptions | str):
if isinstance(color, str):
# If the string doesn't match an Enum member, let it raise an error
color = ColorOptions[color.upper()]
match color:
case ColorOptions.RED:
return 'rgb(1, 0, 0)'
case ColorOptions.BLUE:
return 'rgb(0, 0, 1)'
# Example usage
print(color_something('blue')) # Output: 'rgb(0, 0, 1)'
This works, but it feels inefficient to add this string-to-enum conversion logic to every function that uses an Enum
. One option is to add a helper function, but I'm wondering if there's built-in Enum functionality that I might be missing.
Is there a better way to handle predefined options in a function, where both Enum
members and strings can be accepted without duplicating this logic everywhere?
Edit after @chepner, this limits the options in the IDE instead of str
but then you have to define the options twice (in the Enum class and in the Literal
part of the function).
from typing import Literal
def color_something(color: ColorOptions | Literal['red', 'blue']):
if isinstance(color, str):
# If the string doesn't match an Enum member, let it raise an error
color = ColorOptions[color.upper()]
match color:
case ColorOptions.RED:
return 'rgb(1, 0, 0)'
case ColorOptions.BLUE:
return 'rgb(0, 0, 1)'
# Example usage
print(color_something('blue')) # Output: 'rgb(0, 0, 1)'
The solution is very simple - you need to switch to StrEnum
type, whose members are strings. This way string will be matched to respective enums correctly.
from enum import StrEnum, auto, Enum
class ColorOptions(StrEnum):
RED = auto()
BLUE = auto()
def color_something(color: ColorOptions):
match color:
case ColorOptions.RED:
return 'rgb(1, 0, 0)'
case ColorOptions.BLUE:
return 'rgb(0, 0, 1)'
# Example usage
print(color_something("blue")) # Output: 'rgb(0, 0, 1)'
blue
is lowercase because that is how auto()
works with StrEnum
s. If you want uppercase matches, you can type "RED"
and "BLUE"
in place of auto()
s