如何在Python中为已有枚举添加方法并让IDE识别

问题分析 #

你需要复用包b中枚举B的名称和值，同时在包a中为其添加新行为，但现有两种方法都无法让IDE正常识别新增方法或成员：

• 猴子补丁能运行，但IDE无法识别__call__方法；
• 函数式API动态生成的枚举，IDE无法自动补全成员。

解决方案1：猴子补丁+类型协议（解决方法识别问题） #

通过定义类型协议（Protocol）描述新增方法，再用类型转换让IDE识别枚举的新行为：

from typing import Protocol, cast
from b.enumerations import B
import a

# 定义协议，描述枚举新增的__call__方法
class CallableBProtocol(Protocol):
    def __call__(self, identifier: str) -> a.A: ...

def foo(cls, identifier: str):
    return a.A(namespace=cls.value, identifier=identifier)

# 给枚举打猴子补丁
B.__call__ = foo
# 类型转换，让IDE识别新增方法
B = cast(CallableBProtocol, B)

这样IDE就能识别B的__call__方法，同时保留原枚举的所有成员。

解决方案2：动态生成枚举+类型存根（解决成员补全问题） #

如果需要IDE自动补全枚举成员，可以结合函数式API生成枚举，再添加类型存根文件：

第一步：动态生成枚举（保留原逻辑） #

import enum
from b.enumerations import B as OriginalB
import a

class AEnumItem(str):
    def __call__(self, identifier: str) -> a.A:
        return a.A(namespace=str(self), identifier=identifier)

class AEnumBase(AEnumItem, enum.Enum):
    def __str__(self) -> str:
        return self.value

# 基于原枚举生成新枚举
B = AEnumBase('B', [(item.name, AEnumItem(item.value)) for item in OriginalB])

第二步：添加类型存根文件 #

在包a的目录下创建enumerations.pyi（假设代码在a/enumerations.py），显式声明枚举的成员和方法：

# a/enumerations.pyi
import enum
import a

class B(enum.Enum):
    # 复制原枚举B的所有成员，例如：
    MEMBER_ONE = "member_one_value"
    MEMBER_TWO = "member_two_value"
    # 所有原枚举成员都需在此列出

    def __call__(self, identifier: str) -> a.A: ...

IDE会读取类型存根文件，从而自动补全枚举成员和新增方法。

解决方案3：静态枚举子类+成员复制（折中方案） #

如果不想用类型存根，可以创建新枚举类，手动复制原枚举成员并添加新方法：

import enum
from b.enumerations import B as OriginalB
import a

class B(enum.Enum):
    def __call__(self, identifier: str) -> a.A:
        return a.A(namespace=self.value, identifier=identifier)

# 复制原枚举的所有成员
for name, member in OriginalB.__members__.items():
    setattr(B, name, member.value)

这种方法无需额外文件，IDE能识别新增的__call__方法；缺点是原枚举成员较多时需维护复制代码，可通过脚本自动生成这部分内容。

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