枚举作为函数中的选项，还添加字符串选项

Question

定义函数时，有多种方法可以将输入限制为一组预定义选项。我认为使用

Enum

对象来实现这一点是有意义的。例如：

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

但是，以下调用不起作用：

color_something('BLUE')

为了解决这个问题，我修改了函数以接受

Enum

成员和字符串表示形式，如下所示：

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

这可行，但将这种字符串到枚举转换逻辑添加到每个使用

Enum

的函数感觉效率很低。一种选择是添加辅助函数，但我想知道是否有我可能缺少的内置枚举功能。

是否有更好的方法来处理函数中的预定义选项，其中

Enum

成员和字符串都可以被接受，而无需在各处重复此逻辑？

@chepner 之后的

Edit，这会限制 IDE 中的选项，而不是

str

，但随后您必须定义选项两次（在 Enum 类中和函数的

Literal

部分中）。

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

Answer 1

解决方案很简单——你需要切换到

StrEnum

类型，其成员是字符串。这样字符串将正确匹配相应的枚举。

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

是小写，因为这就是

auto()

与

StrEnum

的配合方式。如果您想要大写匹配，可以键入

"RED"

和

"BLUE"

代替

auto()

s