如何基于IntEnum创建保留自动补全的字符串Enum?
问题描述
现有如下IntEnum定义:
from enum import IntEnum class TransactionTypeIntEnum(IntEnum): UNKNOWN = 0 ADMINISTRATIVE = 1 PAYMENT = 2 TRADE = 3
需要创建一个字符串Enum,成员名称与上述IntEnum完全一致,值为成员名称本身,形式如下:
from enum import Enum class TransactionType(Enum): UNKNOWN = "UNKNOWN" ADMINISTRATIVE = "ADMINISTRATIVE" PAYMENT = "PAYMENT" TRADE = "TRADE"
当前通过BaseEnum.from_iterable动态生成该字符串Enum:
from enum import Enum from typing import Type, Iterable, TypeVar TBaseEnum = TypeVar("TBaseEnum", bound=BaseEnum, covariant=True) class BaseEnum(Enum): @classmethod def from_iterable(cls: Type[TBaseEnum], name: str, values: Iterable[str]) -> Type[TBaseEnum]: return cls(name, {value: value for value in values}) TransactionType = BaseEnum.from_iterable("TransactionType", [element.name for element in TransactionTypeIntEnum])
但该实现存在问题:动态生成的TransactionType丢失了IDE自动补全功能,无法直接提示成员名称,必须查看TransactionTypeIntEnum才能知晓可用成员。
解决方案
方法一:类型检查分支+动态生成(推荐)
利用typing.TYPE_CHECKING常量,在IDE类型检查阶段提供显式的枚举类定义用于自动补全,运行时则动态从IntEnum生成字符串枚举,既保证补全功能,又避免手动同步两个枚举的成员:
from enum import Enum, IntEnum from typing import TYPE_CHECKING # 定义原始IntEnum class TransactionTypeIntEnum(IntEnum): UNKNOWN = 0 ADMINISTRATIVE = 1 PAYMENT = 2 TRADE = 3 if TYPE_CHECKING: # 仅用于IDE类型检查和自动补全的假枚举类 class TransactionType(Enum): UNKNOWN = "UNKNOWN" ADMINISTRATIVE = "ADMINISTRATIVE" PAYMENT = "PAYMENT" TRADE = "TRADE" else: # 运行时动态生成的真实枚举类,与IntEnum成员完全同步 TransactionType = Enum( "TransactionType", {member.name: member.name for member in TransactionTypeIntEnum} )
方法二:代码生成(显式成员定义)
通过exec动态生成枚举类的代码,让IDE能识别到显式的成员定义,同时保证与IntEnum同步:
from enum import Enum, IntEnum class TransactionTypeIntEnum(IntEnum): UNKNOWN = 0 ADMINISTRATIVE = 1 PAYMENT = 2 TRADE = 3 # 生成字符串枚举的代码并执行 transaction_type_code = f""" class TransactionType(Enum): {chr(10).join([f"{member.name} = '{member.name}'" for member in TransactionTypeIntEnum])} """ exec(transaction_type_code)
方法三:元类+类型提示
自定义元类从IntEnum同步成员,同时结合Literal类型提示让IDE识别成员:
from enum import Enum, IntEnum, EnumMeta from typing import Literal, Type, TypeVar TIntEnum = TypeVar("TIntEnum", bound=IntEnum) class StringEnumMeta(EnumMeta): def __new__(cls, name, bases, namespace, source_int_enum: Type[TIntEnum]): # 从IntEnum提取成员名称和值 enum_members = {member.name: member.name for member in source_int_enum} # 创建枚举类 enum_cls = super().__new__(cls, name, bases, enum_members) # 添加Literal类型提示,让IDE识别成员 enum_cls.__annotations__['__members__'] = Literal[tuple(enum_members.keys())] return enum_cls # 定义IntEnum class TransactionTypeIntEnum(IntEnum): UNKNOWN = 0 ADMINISTRATIVE = 1 PAYMENT = 2 TRADE = 3 # 生成字符串枚举 class TransactionType(Enum, metaclass=StringEnumMeta, source_int_enum=TransactionTypeIntEnum): pass
内容的提问来源于stack exchange,提问作者Henrique Andrade
相关产品推荐
相关产品推荐

