enum as option in function, also add string option

October 21, 2024

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)'

>Solution :

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 StrEnums. If you want uppercase matches, you can type "RED" and "BLUE" in place of auto()s