如何让mypy识别传入str fallback时函数返回str类型？

问题描述 #

我有一个函数，返回str类型或指定的回退值（str或None类型）。当传入str类型的回退值时，函数必然返回str类型，但运行mypy时却触发了类型错误。

最小可复现代码：

def f(x: str, fallback: str | None) -> str | None:
    if x.lower() == "yes":
        return "okay"
    return fallback


f("no", fallback="").lower()

报错信息：

f.py:10: error: Item "None" of "str | None" has no attribute "lower"  [union-attr]

解决方案 #

可以通过两种方式让mypy正确推断函数的返回类型：

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

通过重载明确不同参数类型对应的返回类型，让mypy根据传入的fallback类型匹配对应的签名：

from typing import overload

# 重载1：当fallback为str时，返回str
@overload
def f(x: str, fallback: str) -> str: ...

# 重载2：当fallback为None时，返回str | None
@overload
def f(x: str, fallback: None) -> str | None: ...

# 实际实现函数
def f(x: str, fallback: str | None) -> str | None:
    if x.lower() == "yes":
        return "okay"
    return fallback

# 此时mypy能识别返回值为str，调用lower()无报错
f("no", fallback="").lower()

方案2：使用泛型（TypeVar）

通过泛型让mypy根据fallback的类型动态推导返回类型：

from typing import TypeVar, Union

# 定义类型变量，约束为str或None
T = TypeVar("T", str, None)

def f(x: str, fallback: T) -> Union[str, T]:
    if x.lower() == "yes":
        return "okay"
    return fallback

# 传入str类型的fallback时，Union[str, str]会被简化为str，mypy识别安全
f("no", fallback="").lower()

原理说明 #

原函数的返回类型被固定为str | None，mypy无法区分传入str或None作为fallback时的差异。通过重载或泛型，我们明确告诉mypy：

• 当fallback是str时，函数的返回值必然是str（要么是"okay"，要么是传入的str回退值）
• 当fallback是None时，返回值才是str | None

这样mypy就能正确识别调用lower()的安全性。

内容的提问来源于stack exchange，提问作者Nico Schlömer