在mypy中指明函数类型检查

Indicate that function type checks in mypy

我有一个工厂函数,它接受多个 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。对于 Arg1And2Arg2And3,我们知道相同的事情。我理解为什么 mypy 这样做 - 它实际上 不知道 arg1 不是 None,因为 [=] 中可能会发生各种愚蠢的事情23=] - 但是我有什么办法可以告诉 mypy _none_none 保证其参数的类型,而无需在 dispatch_factory 函数中显式添加 cast(float, arg1)

(此外,如果有人有令人信服的替代方法来避免此问题或更清洁,我很高兴听到。)

在分支中验证类型称为类型保护。这目前不可能,但在 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")

通常,类型保护可以重新表述为 return 值,如果它们是正确的类型,否则就什么都不是。这允许在 if statement/expression 中使用守卫(只有在 returned 时才继续)和赋值表达式(存储验证值)。

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")