如何使用python输入提示具有相同参数类型但不同返回类型的函数？

我有一个函数，它的返回类型是 tuple[bool, set[int] | str] 如果第 0 项是 True ，则第 1 项是结果 set[int] ，否则第 1 项是一个 str，显示失败的原因。是这样的

def callee(para_a: int) -> tuple[bool, set[int] | str]:
    result = set([1, 2, 3])
    if (something wrong):
        return False, "something wrong"
    return True, result

现在我从其他函数中调用这个函数

from fastapi import FastAPI
from fastapi.responses import JSONResponse
api = FastAPI()

def kernel(para: set[int]):
    return [i for i in para]

@api.post("/test")
def caller(para_a: int):
    res = callee(para_a)
    if res[0] is True:
        return {"result": kernel(res[1])}
    return JSONResponse(status_code=500, content={"fail_msg": res[1]}

没关系 fastapi ，那只是因为我想说它在web中有时很有用。 现在 mypy 会责怪 error: Argument 1 to "kernel" has incompatible type "set[int] | str"; expected "set[int]" [arg-type] ，所以我想让 mypy 知道如果第0项是 True ，那么第1项是结果 set[int] ，否则第1项是str 。我想了想 overload ，所以就写了

from typing import overload, Literal

@overload
def callee(para_a: int) -> tuple[Literal[True], set[int]]:
    ...
@overload
def callee(para_a: int) -> tuple[Literal[False], str]:
    ...

然后 mypy 怪 Overloaded function signature 2 will never be matched: signature 1's parameter type(s) are the same or broader

我想知道的是有没有好的办法解决我的问题？由于在这种情况下我不能使用 overload ，我应该用什么来让 mypy 知道 res[1] 只是a set[int] 但永远不会是a str if res[0] is True ？

你遇到的问题是 mypy 在这种情况下无法推断出函数重载的正确类型。尽管你定义了两种重载，但 mypy 认为第一种重载（返回 tuple[Literal[True], set[int]] ）可以匹配所有情况，因为它和第二种重载的参数类型相同。

为了解决这个问题，你可以使用 自定义类型提示 来更精确地描述函数的返回类型。

1. 使用 TypeVar 和 Union 定义一个自定义类型

from typing import TypeVar, Union, Tuple, Literal

ResultType = TypeVar('ResultType', set[int], str)

def callee(para_a: int) -> Tuple[Literal[True], set[int]] | Tuple[Literal[False], str]:
    result = set([1, 2, 3])
    if (something wrong):
        return False, "something wrong"
    return True, result

这个代码片段定义了一个名为 ResultType 的类型变量，它可以是 set[int] 或 str 。 然后，函数 callee 的返回类型被定义为两种可能类型的联合类型： Tuple[Literal[True], set[int]] 和 Tuple[Literal[False], str] 。

2. 使用类型守卫

在调用 callee 函数后，可以使用类型守卫来帮助 mypy 理解返回类型的具体情况：

from fastapi import FastAPI
from fastapi.responses import JSONResponse

api = FastAPI()

def kernel(para: set[int]):
    return [i for i in para]

@api.post("/test")
def caller(para_a: int):
    res = callee(para_a)
    if res[0]:  # 类型守卫：如果 res[0] 为 True
        return {"result": kernel(res[1])}  # mypy 知道 res[1] 是 set[int]
    else:
        return JSONResponse(status_code=500, content={"fail_msg": res[1]})  # mypy 知道 res[1] 是 str

在这个例子中， if res[0] 的条件判断充当了类型守卫。当 res[0] 为 True 时， mypy 就会知道 res[1] 的类型是 set[int] ，反之则为 str 。这样就可以避免 mypy 报错，并且代码更加清晰易懂。

通过使用自定义类型提示和类型守卫，你可以更精确地描述函数的行为，并帮助 mypy 进行更准确的类型检查，从而提高代码的可靠性和可维护性。