从函数 return 类型中缩小类型

Question

我有一个函数可能 return 两种不同类型的值，例如Union[str, int]。 returned 值的类型可以从函数参数中确定。所以在每次调用时，我都知道应该 returned 哪种类型。但是，mypy 抱怨 Union[str, int] 类型不能分配给更严格的类型，例如str。有没有办法允许更严格的类型分配，而不完全关闭类型检查（例如使用 # type: ignore 评论）？

例如，如果我有这个功能：

from typing import Union


def change_value(action: str, value: int) -> Union[str, int]
    if action == "stringify":
        return str(value)
    return value

stringified_value: str
stringified_value = change_value("stringify", 10)

Mypy 抱怨这个错误：

Expression of type "str | int" cannot be assigned to declared type "str"

我想通过某种方式告诉 mypy 来消除 mypy 错误，在这种情况下缩小类型总是正确的。

Answer 1

可以使用重载来完成：

from typing import Literal, Union, overload

@overload
def change_value(action: Literal['stringify'], value: int) -> str: ...
@overload
def change_value(action: Literal['whatever_else'], value: int) -> int: ...

def change_value(action: str, value: int) -> Union[str, int]:
    if action == "stringify":
        return str(value)
    return value

stringified_value: str = change_value("stringify", 10)

playground

这种方法还有一个额外的好处：你明确地告诉 stringify 和 whatever_else 是两个唯一可能的操作，并且像 change_value('stingify', 2) 这样的调用将被类型检查器拒绝（注意到打字错误？如果没有，否则很难调试。

从函数 return 类型中缩小类型

Narrowing down a type from function return types

python-3.x

mypy