使用 mypy 处理条件逻辑 + 哨兵值

Question

我有一个大致如下所示的函数：

import datetime
from typing import Union

class Sentinel(object): pass
sentinel = Sentinel()

def func(
    dt: datetime.datetime,
    as_tz: Union[datetime.tzinfo, None, Sentinel] = sentinel,
) -> str:

    if as_tz is not sentinel:
        # Never reached if as_tz has wrong type (Sentinel)
        dt = dt.astimezone(as_tz)
    # ...
    # do other meaningful stuff
    # ...
    return "foo"

sentinel此处使用该值是因为None它已经是 的有效参数.astimezone()，因此目的是正确识别用户根本不想调用的.astimezone()情况。

但是，mypy抱怨这种模式：

错误：“datetime”的“astimezone”的参数 1 具有不兼容的类型“Union[tzinfo, None, Sentinel]”；预期“可选[tzinfo]”

似乎这是因为datetime存根（理所当然）使用：

def astimezone(self, tz: Optional[_tzinfo] = ...) -> datetime: ...

但是，有没有办法让 mypy 知道由于检查而sentinel永远不会传递.astimezone()该if值？或者这只是需要一个# type: ignore没有更清洁的方法？

---

另一个例子：

from typing import Optional
import requests


def func(session: Optional[requests.Session] = None):
    new_session_made = session is None
    if new_session_made:
        session = requests.Session()
    try:
        session.request("GET", "https://a.b.c.d.com/foo")
        # ...
    finally:
        if new_session_made:
            session.close()

第二个，和第一个一样，是“运行时安全的”（因为没有更好的术语）：AttributeErrorfrom 调用None.request()并且None.close()不会被到达或评估。但是，mypy 仍然抱怨：

mypytest.py:9: error: Item "None" of "Optional[Session]" has no attribute "request"
mypytest.py:13: error: Item "None" of "Optional[Session]" has no attribute "close"

我应该在这里做些不同的事情吗？

Answer 1

根据我的经验，最好的解决方案是使用enum.Enum.

要求

一个好的哨兵模式有 3 个属性：

1. 拥有不会被误认为其他值的明确类型/值。例如object()
2. 可以使用描述性常量引用
3. 可以简洁地测试，使用is和is not

解决方案

enum.Enum 由 mypy 特别处理，因此它是我发现的唯一可以满足所有这三个要求并在 mypy 中正确验证的解决方案。

import datetime
import enum
from typing import Union

class Sentinel(enum.Enum):
    SKIP_TZ = object()

def func(
    dt: datetime.datetime,
    as_tz: Union[datetime.tzinfo, None, Sentinel] = Sentinel.SKIP_TZ,
) -> str:

    if as_tz is not Sentinel.SKIP_TZ:
        dt = dt.astimezone(as_tz)
    # ...
    # do other meaningful stuff
    # ...
    return "foo"

变化

这个解决方案还有一些其他有趣的特性。

可重用的 Sentinel 对象

sentinel.py

import enum
class Sentinel(enum.Enum):
    sentinel = object()

main.py

import datetime
from sentinel import Sentinel
from typing import Union

SKIP_TZ = Sentinel.sentinel

def func(
    dt: datetime.datetime,
    as_tz: Union[datetime.tzinfo, None, Sentinel] = SKIP_TZ,
) -> str:

    if as_tz is not SKIP_TZ:
        dt = dt.astimezone(as_tz)
    # ...
    # do other meaningful stuff
    # ...
    return "foo"

请注意，由于Sentinel.sentinel始终提供相同的object实例，因此永远不应使用两个可重用的哨兵相同的上下文。

使用限制哨兵值 Literal

替换Sentinel为Literal[Sentinel.SKIP_TZ]]使您的函数签名更加清晰，尽管它无可否认是多余的，因为只有一个枚举值。

import datetime
import enum
from typing import Union
from typing_extensions import Literal

class Sentinel(enum.Enum):
    SKIP_TZ = object()

def func(
    dt: datetime.datetime,
    as_tz: Union[datetime.tzinfo, None, Literal[Sentinel.SKIP_TZ]] = Sentinel.SKIP_TZ,
) -> str:

    if as_tz is not Sentinel.SKIP_TZ:
        dt = dt.astimezone(as_tz)
    # ...
    # do other meaningful stuff
    # ...
    return "foo"

func(datetime.datetime.now(), as_tz=Sentinel.SKIP_TZ)

不符合我要求的解决方案

自定义哨兵类

import datetime
from typing import Union

class SentinelType:
    pass

SKIP_TZ = SentinelType()


def func(
    dt: datetime.datetime,
    as_tz: Union[datetime.tzinfo, None, SentinelType] = SKIP_TZ,
) -> str:

    if not isinstance(dt, SentinelType):
        dt = dt.astimezone(as_tz)
    # ...
    # do other meaningful stuff
    # ...
    return "foo"

虽然这有效，但使用isinstance(dt, SentinelType)不符合要求 3（“使用is”），因此也不符合要求 2（“使用命名常量”）。为清楚起见，我希望能够使用if dt is not SKIP_TZ.

目的 Literal

Literal 不适用于任意值（尽管它适用于枚举。见上文。）

import datetime
from typing import Union
from typing_extensions import Literal

SKIP_TZ = object()

def func(
    dt: datetime.datetime,
    as_tz: Union[datetime.tzinfo, None, Literal[SKIP_TZ]] = SKIP_TZ,
) -> str:

    if dt is SKIP_TZ:
        dt = dt.astimezone(as_tz)
    # ...
    # do other meaningful stuff
    # ...
    return "foo"

产生以下 mypy 错误：

error: Parameter 1 of Literal[...] is invalid
error: Variable "sentinel.SKIP_TZ" is not valid as a type

细绳 Literal

在这次尝试中，我使用了字符串文字而不是对象：

import datetime
from typing import Union
from typing_extensions import Literal


def func(
    dt: datetime.datetime,
    as_tz: Union[datetime.tzinfo, None, Literal['SKIP_TZ']] = 'SKIP_TZ',
) -> str:

    if as_tz is not 'SKIP_TZ':
        dt = dt.astimezone(as_tz)
    # ...
    # do other meaningful stuff
    # ...
    return "foo"

func(datetime.datetime.now(), as_tz='SKIP_TZ')

即使这行得通，它在需求 1 上也会很弱。

但它没有通过 mypy. 它产生错误：

error: Argument 1 to "astimezone" of "datetime" has incompatible type "Union[tzinfo, None, Literal['SKIP_TZ']]"; expected "Optional[tzinfo]"

Answer 2

您可以使用显式cast：

    from typing import cast
    ... 
    if as_tz is not sentinel:
        # Never reached if as_tz has wrong type (Sentinel)
        as_tz = cast(datetime.tzinfo, as_tz)
        dt = dt.astimezone(as_tz)

和

    new_session_made = session is None
    session = cast(requests.Session, session)

您也可以使用 an assert（尽管这是实际的运行时检查，而cast更明确地说是无操作）：

        assert isinstance(as_tz, datetime.tzinfo)
        dt = dt.astimezone(as_tz)

和

    new_session_made = session is None
    assert session is not None