如何为接收Union类型、返回对应子类型的函数添加类型标注？

原函数与需求 #

原函数实现如下：

def add_one(value):
    match value:
        case int():
            return value + 1
        case str():
            return value + " and one more"
        case _:
            raise TypeError()

需求：

• 函数可接收int（或其子类）或str（或其子类）作为参数
• 传入int或其子类时返回int
• 传入str或其子类时返回str
• 支持传入Union[int, str]类型的参数，返回对应的Union[int, str]

失败的尝试分析 #

尝试1：宽泛的Union标注 #

def add_one(value: int | str) -> int | str: ...

问题：类型检查器无法关联输入与输出类型，会误判传入int可能返回str，无法保证类型一致性。

尝试2：简单泛型标注 #

def add_one[T: int | str](value: T) -> T: ...

问题：函数并非返回与参数完全相同的类型——传入IntEnum（int子类）会返回int而非IntEnum，传入StrEnum（str子类）会返回str而非StrEnum，标注与实际逻辑不符。

尝试3：带约束的泛型标注 #

def add_one[T: (int, str)](value: T) -> T: ...

问题：无法处理Union[int, str]类型的参数，调用时会触发类型检查错误。

尝试4：多重载标注（当前最优） #

@overload
def add_one(value: int) -> int: ...

@overload
def add_one(value: str) -> str: ...

def add_one(value: int | str) -> int | str:
    # 实际实现
    match value:
        case int():
            return value + 1
        case str():
            return value + " and one more"
        case _:
            raise TypeError()

问题：需要为每个支持的类型手动编写重载签名，当类型数量达到7-8种及以上时，代码冗余度会显著提升。

更优解决方案：带边界的TypeVar+重载 #

针对子类返回父类类型的需求，我们可以用带**边界（bound）**的TypeVar来覆盖子类场景，同时保留重载的精确性，且无需为每个子类单独编写签名：

from typing import TypeVar, overload, Union

# 定义绑定到int及其子类的类型变量
TInt = TypeVar("TInt", bound=int)
# 定义绑定到str及其子类的类型变量
TStr = TypeVar("TStr", bound=str)

@overload
def add_one(value: TInt) -> int: ...
@overload
def add_one(value: TStr) -> str: ...
@overload
def add_one(value: Union[TInt, TStr]) -> Union[int, str]: ...

def add_one(value):
    match value:
        case int():
            return value + 1
        case str():
            return value + " and one more"
        case _:
            raise TypeError()

方案优势 #

1. 覆盖子类场景：TInt接受int及其所有子类（如IntEnum），标注明确返回int而非子类类型，符合函数实际逻辑；TStr同理。
2. 支持Union参数：第三个重载处理Union[TInt, TStr]类型的输入，返回对应的Union[int, str]，解决了尝试3的问题。
3. 减少冗余：无需为每个子类单独编写重载，仅需针对基础类型的边界泛型编写签名，即使扩展到7-8种基础类型，也只需对应数量的重载+一个Union重载，比为每个子类写签名简洁得多。

关于“自动匹配Union子类型”的说明 #

Python的静态类型系统目前不支持直接定义“给定Union类型，自动匹配每个子类型并映射返回类型”的泛型逻辑，因为这需要关联类型（Associated Types）的支持，而当前typing模块尚未提供该特性。因此，带边界的TypeVar+重载是当前最简洁且可靠的实现方式。

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