我可以在函数内对一个影响原始参数的Mypy断言吗？

Question

假设我有一个简单的验证函数：

def is_valid_build_target(target: Any, throw=False) -> bool:
    target = str(target)
    allowed_targets = ["dev", "prod"]
    is_allowed = target.lower() in ALLOWED_TARGETS
    if not is_allowed and throw:
        raise ValueError(
            f"Invalid target '{target}'. Must be one of: {ALLOWED_TARGETS}"
        )

    assert target is not None
    return is_allowed

但是如果我想使用这个函数，Mypy（一个检查代码的工具）不会把那个断言传递回上层（可能是因为 target 是一个基本数据类型，所以 Python 会把这个变量限制在函数内部）：

Target = Literal["dev", "prod"]
target: Target | None = cast(Target | None, os.getenv("APP_TARGET", None))

if not is_valid_build_target(target):
    raise ValueError(f"Invalid target, I could have used throw=True, but I wanted a custom error message")

# Mypy still thinks `target` could be None

如果我把函数的逻辑写在一起，或者在后面使用 assert，那就可以正常工作了，但这样我就没有把运行时的验证逻辑放在一个独立的函数里：

if not is_valid_build_target(target):
    raise ValueError(f"Invalid target...")
assert target is not None

# happy Mypy

有没有什么办法可以在验证函数内部进行验证，同时让 Mypy 感到满意呢？

Answer 1

我建议你的验证函数可以做两件事：

1. 返回一个有效的值，或者
2. 抛出一个异常

我还会先检查一下 os.getenv 是否返回 None，因为变量根本不存在是一个不同的问题（虽然解决方法可能相同），而不是有一个无效的字符串值。

import typing
import os


Target = typing.Literal["dev", "prod"]
ALLOWED_TARGETS = typing.get_args(Target)

def valid_build_target(target: str) -> Target:
    lower_target = target.lower()
    if lower_target in ALLOWED_TARGETS:
        # It's a shame type narrowing doesn't work here.
        return typing.cast(Target, lower_target)
    else:
        raise ValueError(
            f"Invalid target '{target}'. Must be one of: {ALLOWED_TARGETS}"
        )

possible_target = os.getenv("APP_TARGET")
if possible_target is None:
    raise ValueError("APP_TARGET not defined")

target: Target = valid_build_target(possible_target)
   

尽管 possible_target 最开始的类型是 str | None，但在你调用 valid_build_target 的时候，它的类型已经变成了 str。

$ python3 tmp.py  # Raises
$ APP_TARGET= python3 tmp.py   # Raises
$ APP_TARGET=foo python3 tmp.py  # Raises
$ APP_TARGET=DEV python3 tmp.py  # OK
$ APP_TARGET=dev python3 tmp.py  # OK