如何为Python函数指定输入与返回值关联的条件类型注解？

当然可以实现这种输入类型和返回类型的对应关系，Python的类型系统支持两种常用方式来做到这一点：

方法一：使用@overload装饰器 #

这是最直观的方案，针对不同的输入类型分别定义函数签名，让类型检查器明确识别输入输出的对应关系：

from typing import overload, str, int, float, list

# 单个字符串/整数输入时的签名
@overload
def get(ids: str | int) -> float:
    ...

# 列表类型输入时的签名
@overload
def get(ids: list[str] | list[int]) -> list[float]:
    ...

# 实际实现的函数体
def get(ids: str | list[str] | int | list[int]) -> float | list[float]:
    if isinstance(ids, list):
        return [float(item) for item in ids]
    return float(ids)

当你传入单个str或int时，类型检查器（如mypy、pyright）会自动推断返回值为float；传入list[str]或list[int]时，会推断返回list[float]。

方法二：结合TypeVar与泛型 #

如果需要更灵活的泛型约束，可以用TypeVar定义关联类型：

from typing import TypeVar, float, list

# 定义允许的单个元素类型
SingleType = TypeVar('SingleType', str, int)
# 定义允许的列表类型
ListType = TypeVar('ListType', list[str], list[int])

def get(ids: SingleType | ListType) -> float | list[float]:
    if isinstance(ids, list):
        return [float(item) for item in ids]
    return float(ids)

这种方式同样能让类型检查器识别输入输出的对应关系，适合需要扩展更多类型场景的情况。

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