如何让Python函数同时支持Enum成员与字符串参数输入？

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

# 示例调用
print(color_something(ColorOptions.BLUE))  # 输出: 'rgb(0, 0, 1)'

但这种实现存在局限：直接传入字符串调用会触发类型错误，比如：

color_something('BLUE')  # 运行报错

为了让函数同时支持Enum成员和字符串，我们可以在函数内添加字符串转Enum的逻辑：

def color_something(color: ColorOptions | str):
    if isinstance(color, str):
        # 字符串不匹配时自动抛出错误
        color = ColorOptions[color.upper()]
    match color:
        case ColorOptions.RED:
            return 'rgb(1, 0, 0)'        
        case ColorOptions.BLUE:
            return 'rgb(0, 0, 1)'

# 示例调用
print(color_something('blue'))  # 输出: 'rgb(0, 0, 1)'

这种方法虽然可行，但如果多个函数都需要处理相同的Enum，重复编写转换逻辑会非常冗余。以下是几种更优的解决方案，彻底避免重复代码：

方法一：给Enum类封装字符串转换逻辑 #

在Enum类中添加类方法，统一处理字符串到Enum成员的转换，所有需要转换的地方直接调用该方法即可：

from enum import Enum, auto

class ColorOptions(Enum):
    RED = auto()
    BLUE = auto()

    @classmethod
    def from_str(cls, color_str: str):
        return cls[color_str.upper()]

def color_something(color: ColorOptions | str):
    if isinstance(color, str):
        color = ColorOptions.from_str(color)
    match color:
        case ColorOptions.RED:
            return 'rgb(1, 0, 0)'        
        case ColorOptions.BLUE:
            return 'rgb(0, 0, 1)'

方法二：用装饰器统一处理参数转换 #

编写通用装饰器，自动将指定参数中的字符串转换为对应Enum类型，彻底消除重复的转换代码：

from enum import Enum, auto
from functools import wraps

def enum_str_converter(enum_cls, param_name):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # 处理关键字参数
            if param_name in kwargs and isinstance(kwargs[param_name], str):
                kwargs[param_name] = enum_cls[kwargs[param_name].upper()]
            # 处理位置参数
            new_args = list(args)
            for idx, arg in enumerate(new_args):
                param_keys = list(func.__annotations__.keys())
                if idx < len(param_keys) and param_keys[idx] == param_name and isinstance(arg, str):
                    new_args[idx] = enum_cls[arg.upper()]
            return func(*new_args, **kwargs)
        return wrapper
    return decorator

class ColorOptions(Enum):
    RED = auto()
    BLUE = auto()

@enum_str_converter(ColorOptions, 'color')
def color_something(color: ColorOptions | str):
    match color:
        case ColorOptions.RED:
            return 'rgb(1, 0, 0)'        
        case ColorOptions.BLUE:
            return 'rgb(0, 0, 1)'

# 两种调用方式均有效
print(color_something(ColorOptions.RED))
print(color_something('blue'))

方法三：自动生成Literal类型提示，避免重复定义选项 #

针对使用Literal时需要手动重复字符串选项的问题，可以通过Enum的内置属性自动生成Literal类型，无需手动维护字符串列表：

from enum import Enum, auto
from typing import Literal, TypeVar

T = TypeVar('T', bound=Enum)

def enum_literal(enum_cls: type[T]) -> type:
    return Literal[tuple(enum_cls.__members__.keys())]

class ColorOptions(Enum):
    RED = auto()
    BLUE = auto()

# 自动生成Literal['RED', 'BLUE']类型
ColorLiteral = enum_literal(ColorOptions)

def color_something(color: ColorOptions | ColorLiteral):
    if isinstance(color, str):
        color = ColorOptions[color.upper()]
    match color:
        case ColorOptions.RED:
            return 'rgb(1, 0, 0)'        
        case ColorOptions.BLUE:
            return 'rgb(0, 0, 1)'

# IDE会自动提示可选的字符串选项
print(color_something('blue'))

内容的提问来源于stack exchange，提问作者3dSpatialUser