如何为Python类中的单个方法添加正确的泛型类型标注？

问题背景 #

当前代码中Container.get_by_type方法返回List[BaseClass]，但实际根据传入的Types枚举，它应该返回对应子类（如ClassA、ClassB）的列表，导致类型检查工具抛出“类型不兼容”的错误。

解决方案 #

通过添加类型映射、自定义绑定子类的TypeVar，配合类型守卫或类型转换，让类型检查器能正确推断get_by_type的返回类型。

完整修改后的代码

from dataclasses import dataclass
from enum import Enum, auto
from typing import List, TypeVar, Generic, Type, cast, TypeGuard

T = TypeVar("T")
# 绑定到BaseClass子类的TypeVar，用于标注具体返回类型
BC = TypeVar("BC", bound=BaseClass)

@dataclass
class BaseClass(Generic[T]):
    def validate(self, value: T):
        raise NotImplementedError

class Types(Enum):
    A = auto()
    B = auto()

    # 枚举值与对应子类的映射表
    _TYPE_MAP: dict["Types", Type[BaseClass]] = {}

    @staticmethod
    def from_instance(instance: BaseClass) -> "Types":
        if isinstance(instance, ClassA):
            return Types.A
        if isinstance(instance, ClassB):
            return Types.B
        raise ValueError("Not supported")

    @classmethod
    def get_class(cls, type_enum: "Types") -> Type[BC]:
        return cast(Type[BC], cls._TYPE_MAP[type_enum])

@dataclass
class ClassA(BaseClass[str]):
    def validate(self, value: str):
        pass

@dataclass
class ClassB(BaseClass[int]):
    def validate(self, value: int):
        pass

# 注册子类到枚举映射
Types._TYPE_MAP[Types.A] = ClassA
Types._TYPE_MAP[Types.B] = ClassB

# 可选：定义类型守卫函数，让类型推断更严谨
def is_target_type(instance: BaseClass, type_enum: Types) -> TypeGuard[BC]:
    return isinstance(instance, Types.get_class(type_enum))

@dataclass
class Container:
    instances: List[BaseClass]

    def get_by_type(self, type: Types) -> List[BC]:
        # 使用类型守卫过滤，无需额外类型转换
        return [instance for instance in self.instances if is_target_type(instance, type)]

# 测试代码
def function(fields_from_class_a: List[ClassA]):
    print(fields_from_class_a)

container = Container(instances=[ClassA(), ClassB()])
fields = container.get_by_type(Types.A)
function(fields)  # 类型检查器会正确识别fields为List[ClassA]

关键修改点说明 #

1. 自定义TypeVar：新增BC并绑定到BaseClass子类，让方法可以返回具体的子类类型，而非模糊的父类类型。
2. 枚举类型映射：Types._TYPE_MAP建立枚举值与子类的关联，避免硬编码重复的类型判断逻辑。
3. 类型守卫/类型转换：通过TypeGuard或cast告知类型检查器，过滤后的实例属于目标子类，消除类型不兼容的报错。

内容的提问来源于stack exchange，提问作者Antonio Gamiz Delgado