Python函数重载中区分同质与异质元组的类型注解问题

问题根源 #

你遇到的MyPy报错，本质是重载签名的顺序和范围逻辑错误：如果先定义泛化的*items: Base重载，再定义具体的*items: Z重载，MyPy会优先匹配泛化签名，导致具体签名永远不会被触发。正确的做法是从具体到泛化排列重载，同时用TypeVar明确区分同质/异质输入的类型约束。

完整代码实现 #

from abc import ABC, abstractmethod
from typing import TypeVar, overload, FrozenSet, Union

# 抽象基类与子类
class Base(ABC):
    @abstractmethod
    def do_something(self) -> None:
        pass

class A(Base):
    def do_something(self) -> None:
        pass

class B(Base):
    def do_something(self) -> None:
        pass

class Z(Base):
    def do_something(self) -> None:
        pass

# Product接口及实现
class Product(ABC):
    @abstractmethod
    def get_items(self) -> FrozenSet[Base]:
        pass

class MixedProduct(Product):
    def __init__(self, items: FrozenSet[Base]) -> None:
        self._items = items
    def get_items(self) -> FrozenSet[Base]:
        return self._items

class ZProduct(Product):
    def __init__(self, items: FrozenSet[Z]) -> None:
        self._items = items
    def get_items(self) -> FrozenSet[Z]:
        return self._items

# 定义绑定Base的TypeVar，用于泛化场景的类型注解
T = TypeVar('T', bound=Base)

# 重载1：最具体场景——所有输入都是Z类型
@overload
def make_product(*items: Z) -> ZProduct: ...

# 重载2：所有其他场景——输入为任意Base子类（混合类型或非Z同质类型）
@overload
def make_product(*items: T) -> MixedProduct: ...

# 工厂函数实现
def make_product(*items: Base) -> Union[ZProduct, MixedProduct]:
    def is_all_z(items: tuple[Base, ...]) -> bool:
        return all(isinstance(item, Z) for item in items)
    
    if is_all_z(items):
        return ZProduct(frozenset(items))
    return MixedProduct(frozenset(items))

关键说明 #

1. 重载顺序：必须先写最具体的*items: Z重载，再写泛化的*items: T重载，MyPy会从上到下匹配第一个符合条件的签名。
2. TypeVar约束：用TypeVar('T', bound=Base)表示任意Base的子类，覆盖所有非全Z的输入情况（包括混合类型、全A/全B等同质非Z类型）。
3. 类型推断辅助：内部的is_all_z函数会让MyPy自动推断：当函数返回True时，输入元组的类型为tuple[Z, ...]，从而绑定ZProduct的返回类型。

验证示例 #

以下调用场景的类型推断均能被MyPy正确识别：

# 全Z输入 → 推断返回ZProduct
z1, z2 = Z(), Z()
zp = make_product(z1, z2)
reveal_type(zp)  # MyPy输出: ZProduct

# 混合类型输入 → 推断返回MixedProduct
a = A()
mp1 = make_product(z1, a)
reveal_type(mp1)  # MyPy输出: MixedProduct

# 全A输入 → 推断返回MixedProduct
a1, a2 = A(), A()
mp2 = make_product(a1, a2)
reveal_type(mp2)  # MyPy输出: MixedProduct

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