如何修复Python项目中类Enum风格SimpleEnum的mypy返回值类型不兼容错误

我们项目未使用标准库Enum，而是自定义了SimpleEnum，定义如下：

class SimpleEnum:
    @classmethod
    def values(cls):
        return {
            getattr(cls, name)
            for name in dir(cls)
            if not name.startswith("_") and not callable(getattr(cls, name))
        }

编写如下代码时：

class MySimpleEnum(SimpleEnum):
    completed = "completed"
    parse = "parse"


def foo() -> MySimpleEnum:
    return MySimpleEnum.parse

会触发mypy错误：

error: Incompatible return value type (got "str", expected "MySimpleEnum") [return-value]

项目规模达数十万行，无法直接切换到标准Enum，需在不改变SimpleEnum原有行为（已被数百处使用）的前提下修改它，以支持上述类型提示。

可以通过给SimpleEnum添加泛型类型注解和类型提示相关的钩子方法，让mypy识别子类成员的类型，同时完全保留原有运行时逻辑。

方案一：泛型子类显式指定类型 #

修改SimpleEnum为泛型类，添加__class_getitem__方法用于类型提示：

from typing import TypeVar, Generic

T = TypeVar('T', str, int, bool)

class SimpleEnum(Generic[T]):
    @classmethod
    def values(cls) -> set[T]:
        return {
            getattr(cls, name)
            for name in dir(cls)
            if not name.startswith("_") and not callable(getattr(cls, name))
        }
    
    @classmethod
    def __class_getitem__(cls, item) -> "Type[T]":
        # 运行时返回自身，仅为类型检查器提供信息
        return cls

定义子类时指定泛型参数，类型注解中同步使用：

class MySimpleEnum(SimpleEnum[str]):
    completed = "completed"
    parse = "parse"


def foo() -> MySimpleEnum[str]:
    return MySimpleEnum.parse

此方案完全不影响原有代码的运行行为，values()方法依然返回成员值的集合，子类成员保持原始类型，仅为mypy补充类型信息。

方案二：兼容旧注解的隐式类型推断 #

如果不想修改现有代码的类型注解，可以添加__getattr__的类型提示，让mypy自动推断成员类型：

from typing import TypeVar, Generic, Type

T = TypeVar('T', str, int, bool)

class SimpleEnum(Generic[T]):
    @classmethod
    def values(cls) -> set[T]:
        return {
            getattr(cls, name)
            for name in dir(cls)
            if not name.startswith("_") and not callable(getattr(cls, name))
        }
    
    @classmethod
    def __class_getitem__(cls, item: Type[T]) -> Type[T]:
        return cls
    
    @classmethod
    def __getattr__(cls, name: str) -> T:
        # 仅用于类型提示，运行时由Python原生属性查找处理
        raise AttributeError(f"{cls.__name__} has no attribute {name}")

这样原有代码中的def foo() -> MySimpleEnum:注解可以保留，mypy会自动将MySimpleEnum.parse识别为str类型，通过类型检查。

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