在 mypy 中表示函数类型检查

Question

我有一个工厂函数，它接受多个 Optional 参数并根据哪些参数是 None 而不是创建对象。我还有一个函数可以检查它的所有参数是否都是None，还有一个函数可以检查它的所有参数是否不是None。我用这些来做检查。这是 MWE：

from typing import Optional

class Arg1:
    def __init__(self, arg1: float):
        pass

class Arg1And2:
    def __init__(self, arg1: float, arg2: float):
        pass

class Arg2And3:
    def __init__(self, arg2: float, arg3: float):
        pass

def _all_none(*args) -> bool:
    return all(a is None for a in args)

def _none_none(*args) -> bool:
    return all(a is not None for a in args)

def dispatch_factory(
    arg1: Optional[float] = None,
    arg2: Optional[float] = None,
    arg3: Optional[float] = None
):
    if _all_none(arg2, arg3) and _none_none(arg1):
        return Arg1(arg1)  # <-- This throws a mypy error
    elif _all_none(arg3) and _none_none(arg1, arg2):
        return Arg1And2(arg1, arg2)  # <-- And this
    elif _all_none(arg1) and _none_none(arg2, arg3):
        return Arg2And3(arg2, arg3)  # <-- And this
    else:
        raise ValueError("Combination of arguments invalid")

当然，我们知道我构造的Arg1 是有效的，因为我们刚刚检查了arg1 不是None，因此必须是float。对于Arg1And2 和Arg2And3，我们知道同样的事情。我明白为什么 mypy 会这样做 - 它实际上无法知道 arg1 不是 None，因为在 _none_none 中可能会发生各种愚蠢的事情 - 但有什么办法吗让我告诉 mypy _none_none 保证其参数的类型，而不在 dispatch_factory 函数中显式添加 cast(float, arg1)？

（另外，如果有人对此方法有一个令人信服的替代方案，可以避免此问题或更简洁，我很高兴听到。）

Answer 1

在分支中验证类型称为类型保护。这目前是不可能的，但在PEP 647 -- User-Defined Type Guards 中提出。

from typing import TypeGuard, Optional, TypeVar

T = TypeVar('T')

def all_none(args: list[Optional[T]]) -> TypeGuard[List[None]]:
    return all(a is None for a in args)

---

由于代码必须明确地将每个参数传递给守卫，展开检查的工作量大致相同——这适用于当前类型检查。

def dispatch_factory(
    arg1: Optional[float] = None,
    arg2: Optional[float] = None,
    arg3: Optional[float] = None
):
    if arg1 is not None and arg2 is None and arg3 is None:
        return Arg1(arg1)
    elif arg1 is not None and arg2 is not None and arg3 is None:
        return Arg1And2(arg1, arg2)
    elif arg1 is None and arg2 is not None and arg3 is not None:
        return Arg2And3(arg2, arg3)
    else:
        raise ValueError("Combination of arguments invalid")

通常情况下，类型保护可以被重新制定以返回值，如果它们是正确的类型，否则什么都没有。这允许在 if 语句/表达式（仅在返回任何内容时继续）和赋值表达式（以存储验证值）中使用守卫。

def _all_none(*args: Optional[T]) -> List[None]:
    if all(a is None for a in args):  # MyPy does not understand this to narrow `args`, so we need the next ignore
        return list(args)  # type: ignore
    return []

def _none_none(*args: Optional[T]) -> List[T]:
    if all(a is not None for a in args):
        return list(args)  # type: ignore
    return []

def dispatch_factory(
    arg1: Optional[float] = None,
    arg2: Optional[float] = None,
    arg3: Optional[float] = None
):
    if _all_none(arg2, arg3) and (args :=_none_none(arg1)):
        return Arg1(*args)
    elif _all_none(arg3) and (args := _none_none(arg1, arg2)):
        return Arg1And2(*args)
    elif _all_none(arg1) and (args := _none_none(arg1, arg2)):
        return Arg2And3(*args)
    else:
        raise ValueError("Combination of arguments invalid")