python - 使用 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 提示这种模式:

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

这似乎是因为 datetime stub (理所当然地)使用:

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

第二个，和第一个一样，是“运行时安全的”(因为没有更好的术语):AttributeError 来自调用 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. 拥有明确的类型/值，不会被误认为是其他值。例如对象()
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 不适用于任意值(尽管它确实适用于枚举。见上文。)

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

字符串文字

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

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