基于TypeGuard的条件判断中mypy无法收窄泛型类型

你遇到的核心问题是：mypy对绑定到具体类型的TypeVar（如AB = t.TypeVar('AB', A, B)）在联合类型（AB | Iterable[AB]）的收窄上存在局限性，无法自动推断is_a_or_b(arg)为False时arg必然是Iterable[AB]；同时在__setitem__方法中，mypy无法结合重载信息与TypeGuard完成精确类型收窄。以下是针对性的解决方案：

解决方案1：拆分类型守卫，细化分支逻辑 #

针对受约束的TypeVar，拆分通用类型守卫为针对具体类型的守卫，让mypy能逐个识别类型分支，避免对Iterable[AB]这种泛型可迭代类型的直接推断障碍。

修改后的基础示例代码 #

from __future__ import annotations
import typing as t
from collections import abc
from itertools import tee
from typing_extensions import reveal_type

class A:
    ...

class B:
    ...

AB = t.TypeVar('AB', A, B)

T = t.TypeVar('T')

def is_iterable_of(s: abc.Iterable[t.Any], _type: t.Type[T]) -> t.TypeGuard[abc.Iterable[T]]:
    return all(isinstance(x, _type) for x in s)

def is_a(x: t.Any) -> t.TypeGuard[A]:
    return isinstance(x, A)

def is_b(x: t.Any) -> t.TypeGuard[B]:
    return isinstance(x, B)

def is_iterable_of_a(s: abc.Iterable[t.Any]) -> t.TypeGuard[abc.Iterable[A]]:
    return is_iterable_of(s, A)

def is_iterable_of_b(s: abc.Iterable[t.Any]) -> t.TypeGuard[abc.Iterable[B]]:
    return is_iterable_of(s, B)

def fn(arg: AB | abc.Iterable[AB]):
    if is_a(arg) or is_b(arg):
        reveal_type(arg)  # 正确收窄为A | B
    elif is_iterable_of_a(arg):
        reveal_type(arg)  # 正确收窄为Iterable[A]
    elif is_iterable_of_b(arg):
        reveal_type(arg)  # 正确收窄为Iterable[B]
    else:
        raise TypeError("无效类型")

解决方案2：用typing.cast辅助类型推断 #

在mypy无法自动完成类型收窄的场景下，使用typing.cast显式指定类型，结合断言保证运行时类型安全，同时帮助mypy完成推断。

基础示例else分支修改 #

def fn(arg: AB | abc.Iterable[AB]):
    if is_a_or_b(arg):
        reveal_type(arg)  # A | B
    else:
        # 显式声明arg为Iterable[AB]
        arg_iter = t.cast(abc.Iterable[AB], arg)
        assert is_iterable_of_ab(arg_iter)
        reveal_type(arg_iter)  # Iterable[AB]

TypePureList的__setitem__修改 #

class TypePureList(t.Generic[AB]):
    def __init__(self, items: abc.Iterable[AB]):
        if not isinstance(items, list):
            items = list(items)
        self._items: list[AB] = items

    @t.overload
    def __setitem__(self, index: t.SupportsIndex, value: AB):
        ...

    @t.overload
    def __setitem__(self, index: slice, value: abc.Iterable[AB]):
        ...

    def __setitem__(self, index: t.SupportsIndex | slice, value: AB | abc.Iterable[AB]):
        if len(self._items) == 0:
            raise ValueError

        if isinstance(index, t.SupportsIndex):
            value_single = t.cast(AB, value)
            assert is_a_or_b(value_single)
            assert is_iterable_of_ab([self._items[0], value_single])
            self._items[index] = value_single
        else:
            # 显式声明value为Iterable[AB]
            value_iter = t.cast(abc.Iterable[AB], value)
            value_iter, value1, value2 = tee(value_iter, 3)
            assert is_iterable_of_ab(value1)
            assert is_iterable_of_ab([self._items[0], next(value2)])
            self._items[index] = value_iter

解决方案3：利用泛型容器的已有类型约束 #

对于TypePureList这种泛型容器，可直接依赖泛型的类型约束，确保调用__setitem__时传入的参数类型合法，无需重复检查类型匹配：

class TypePureList(t.Generic[AB]):
    def __init__(self, items: abc.Iterable[AB]):
        if not isinstance(items, list):
            items = list(items)
        self._items: list[AB] = items

    @t.overload
    def __setitem__(self, index: t.SupportsIndex, value: AB):
        ...

    @t.overload
    def __setitem__(self, index: slice, value: abc.Iterable[AB]):
        ...

    def __setitem__(self, index: t.SupportsIndex | slice, value: AB | abc.Iterable[AB]):
        if len(self._items) == 0:
            raise ValueError

        # 直接利用泛型约束，mypy会自动校验类型合法性
        if isinstance(index, t.SupportsIndex):
            self._items[index] = value
        else:
            self._items[index] = value

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