Python Covariance TypeVar问题：如何根据输入类型列表缩小函数返回类型？

核心问题 #

需要实现一个函数，根据传入的允许类型集合缩小返回类型（仅返回A/B实例，排除C），但Pyright对返回值的类型推断报错，不想使用cast强制转换。

错误原因 #

1. Error 1：type(item) in allowed_items是纯运行时检查，类型检查器无法将item的类型从Base关联到Items_co，因此无法确认item符合返回类型约束。
2. Error 2：固定返回A.from_db()会和Items_co的推断类型冲突——如果调用者传入allowed_items={B}，Items_co会被推断为B，而A不能赋值给B。

解决方案 #

通过调整类型检查逻辑和默认返回值生成方式，让类型检查器能正确推断返回类型：

from __future__ import annotations

from collections.abc import Collection
from typing import TypeVar, Union
from typing_extensions import Self


class Base:
    @classmethod
    def from_db(cls: type[Self]) -> Self:
        return cls()

class A(Base): ...
class B(Base): ...
class C(Base): ...

ItemUnions = Union[A, B]
Items_co = TypeVar("Items_co", bound=ItemUnions, covariant=True)

items: Collection[Base] = [A(), B(), C()]

def get_item(allowed_items: Collection[type[Items_co]] = frozenset((A, B))) -> Items_co:
    for item in items:
        # 遍历允许类型，用isinstance做类型守卫，让类型检查器缩小item类型
        for cls in allowed_items:
            if isinstance(item, cls):
                return item
    # 从允许类型集合中取默认类生成实例，确保类型匹配Items_co
    default_cls = next(iter(allowed_items))
    return default_cls.from_db()

关键改进点 #

• 用isinstance(item, cls)替代type(item) in allowed_items：遍历允许类型逐个检查，类型检查器能识别isinstance的类型守卫作用，将item的类型缩小为cls的实例（即Items_co）。
• 动态生成默认返回值：从allowed_items中取第一个类调用from_db()，确保返回值类型始终符合Items_co的约束，避免固定返回某个子类导致的类型不匹配。

验证类型推断 #

调用函数时，类型检查器能正确缩小返回类型：

a = get_item(allowed_items={A, B})  # 推断类型：A | B
b = get_item(allowed_items={A})     # 推断类型：A
c = get_item()                      # 推断类型：A | B

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