如何为函数添加基于输入“形态”的返回类型注解？

但调用另一个只接收单个`str`的函数时，类型检查直接报错：
```python
def bar(s: str):
    """这个函数只接受单个`str`"""
    ...

b = bar(foo(33)) # 类型检查失败：Type "str | list[str]" cannot be assigned to type "str"

检查器给出的警告是：
"Type "str | list[str]" cannot be assigned to type "str""

注：int和str只是示例，实际可以替换成其他类。

怎么给foo加类型注解才能让检查器识别输入和返回的对应关系——比如传单个int就知道返回单个str，从而避免报错？

解决方法 #

方法1：用函数重载（@overload）

这是最直接的方案，给foo定义多个类型签名，明确输入和返回的一一对应关系，类型检查器会根据你实际传的参数匹配对应的签名，精准推断返回类型。

代码示例：

from typing import overload, Union, List

# 先定义「单个int输入→单个str输出」的签名
@overload
def foo(a: int) -> str:
    ...

# 再定义「list[int]输入→list[str]输出」的签名
@overload
def foo(a: List[int]) -> List[str]:
    ...

# 实际实现函数，这里的注解可以用Union兼容两种情况
def foo(a: Union[int, List[int]]) -> Union[str, List[str]]:
    if isinstance(a, list):
        return [str(num) for num in a]
    else:
        return str(a)

这样调用foo(33)时，检查器会匹配第一个签名，确定返回的是str，传给bar就不会报错；传列表时则匹配第二个签名，返回list[str]。

方法2：用泛型（适配任意类型映射）

如果你的函数需要支持任意输入类型T和对应的输出类型U（不局限于int转str），可以用TypeVar定义类型变量，实现更通用的注解：

from typing import TypeVar, List, Union

# 定义输入、输出的类型变量
T = TypeVar('T')
U = TypeVar('U')

# 重载泛型签名
@overload
def foo(a: T) -> U:
    ...

@overload
def foo(a: List[T]) -> List[U]:
    ...

# 实际实现，这里的转换逻辑可以替换成你需要的T→U逻辑
def foo(a: Union[T, List[T]]) -> Union[U, List[U]]:
    if isinstance(a, list):
        return [str(item) for item in a]  # 示例：把每个T转成str（U）
    else:
        return str(a)

为什么原来的Union不行？ #

用Union的话，类型检查器只能知道输入是两种类型之一、返回是两种类型之一，但没法建立「输入单个→返回单个，输入列表→返回列表」的关联关系。哪怕你传的是单个int，它仍然认为返回可能是list[str]，所以传给需要str的bar会报错。而重载刚好能解决这个关联问题。

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