如何根据参数的值更改函数的暗示返回类型

Question

我有一个函数，它通常返回一个它搜索的对象并执行一些其他操作。如果找不到匹配项，它会引发异常。通常，我不在乎它是否找到匹配项，但频率不够高，以至于我会考虑完全删除异常。因此，我做出的妥协是创建一个默认为 True 的参数 raise_on_failure。如果为 False，则返回 None 而不是引用。

def searchAndOperate(f, raise_on_failure: bool = True) -> Optional[Foo]:
    # Search
    result = ...
    if result is None:
        if raise_on_failure: raise ValueError("No results")
        else: return None
    # Operate
    ...
    # Then return the result
    return result

然而，这给我的类型提示带来了一些问题。我想这样做，以便如果 raise_on_failure 为 True，该函数保证返回 Foo，这样它就只是 如果 raise_on_failure 为 False，则可选[Foo]。

我该怎么做？类似于以下代码片段的东西将是最理想的，但我也对其他想法持开放态度。

def searchAndOperate(
    f,
    raise_on_failure: bool = True
) -> Foo if raise_on_failure else Optional[Foo]:
    ...

Answer 1

其他人已经提到，将函数 searchAndOperate 重构为两个单独的函数可能会更好。

然而，为了完整起见，实际上可以使用 @overload 装饰器：

@overload
def searchAndOperate(f, raise_on_failure: Literal[True] = True) -> Foo: ...
@overload
def searchAndOperate(f, raise_on_failure: Literal[False]) -> Optional[Foo]: ...
@overload
def searchAndOperate(f, raise_on_failure: bool = True) -> Optional[Foo]: ...
def searchAndOperate(f, raise_on_failure = True):
    # Search
    result = ...
    if result is None:
        if raise_on_failure: raise ValueError("No results")
        else: return None
    # Operate
    ...
    # Then return the result
    return result

请注意，每个组合都有一个重载，再加上一个通用后备。对于更全面的函数签名，重载列表可能会变得很长。像 mypy 这样的静态类型检查器现在推断以下返回类型：

a = searchAndOperate(f)  # Foo
b = searchAndOperate(f, raise_on_failure=True)  # Foo
c = searchAndOperate(f, raise_on_failure=False)  # Foo | None
d = searchAndOperate(f, raise_on_failure=(10 > 2))  # Foo | None
e = searchAndOperate(f, raise_on_failure=boolarg)  # Foo | None

当不使用关键字参数或文字参数 True 或 False 时，类型检查器正确推断两个专门的返回类型之一过载。当传递布尔表达式或变量时，类型检查器只能推断泛型重载（因为它无法在不运行代码的情况下知道实际值）。

Answer 2

从概念上讲，我认为根据参数而变化的类型提示没有意义。

类型提示 Optional[Foo] 已经封装了 searchAndOperate 能够返回的所有内容：任何调用它的人都知道它将返回 Foo 或无。尝试根据 raise_on_failure 的值“拆分”返回类型会将函数的行为与其实现耦合起来。您应该做的是重构函数本身，使其行为明确。对当前类型提示的不满意源于 searchAndOperate 尝试同时执行太多操作这一事实。

例如，searchAndOperate 可以分为单职责函数 search 和 operate。或者您可以完全放弃 raise_on_failure 参数，并让调用者尝试/排除函数调用。如果它成为您经常做的事情，您可以将其包装为一个单独的函数，例如：

def searchAndOperate(f) -> Foo:
    """Searches and returns Foo. Raises a ValueError if no result is found."""
    result = ...
    if result is None:
        raise ValueError("No result")
    return result

def trySearchAndOperate(f) -> Optional[Foo]:
    """Searches and returns Foo. Returns None if no result is found."""
    try:
        result = searchAndOperate(f)
    except ValueError:
        return None

这样，调用者就已经知道在没有找到结果时会发生什么。您可能认为这有点重复，但这是一种常见模式：'abc'.index(s) 和 'abc'.find(s) 之间的唯一区别> 的区别是，如果在 'abc' 中找不到 s，第一个会引发异常，而第二个则返回默认值 -1在这种情况下。

Answer 3

我不确定这是否可能。

返回类型提示应该告诉您函数返回的内容，无需查看参数的值。

您本质上在此函数中实现了两种不同的行为：

1. 返回 Foo 或崩溃
2. 返回 Foo 或 None

如果您想将其编码为输入提示，您必须考虑这两种情况。这意味着函数是Optional[Foo]。

如果您不希望这样，我建议将此函数拆分为两个不同的函数；一种是加注，另一种是加点：

def search_and_operate(f) -> Optional[Foo]:
    result = ... # some code that either results in a Foo object or None
    return result

def search_and_operate_or_raise(f) -> Foo:
    result = search_and_operate(f)
    if result is None:
        raise ValueError("No result")
    return result