字符串转字面量时出现不兼容类型错误

Question

运行下面这段代码时，使用了mypy工具：

from typing import Literal, Final

def extract_literal(d2: Literal["b", "c"]) -> str:
    if d2 == "b":
        return "BA"
    if d2 == "c":
        return "BC"
    
def model(d2_name: str = "b-123") -> None:
    if d2_name[0] not in ["b", "c"]:
        raise AssertionError
    d2: Final = d2_name[0]
    print(extract_literal(d2))

结果是：

typing_test.py:17: error: Argument 1 to "extract_literal" has incompatible type "str";
expected "Literal['b', 'c']"  [arg-type]
        print(extract_literal(d2))
                              ^~
Found 1 error in 1 file (checked 1 source file)

为了让你更明白，d2_name这个变量肯定是要么是"b-number"，要么是"c-number"。根据它的第一个字母，我想输出不同的消息。

Python version: 3.11
mypy version: 1.9.0

那么，想要让mypy通过，需要做什么改动呢？

Answer 1

为了让大家明白，d2_name的值一定是“b-number”或者“c-number”中的一个。

你是怎么知道的呢？

如果你是直接在程序的其他地方写死了d2_name的值，那就一开始就用你期望的格式来写。可以参考一下mypy playground。

Name: TypeAlias = tuple[Literal["b", "c"], int]

EXAMPLE_D2_NAME_1: Name = ("b", 123)
EXAMPLE_D2_NAME_2: Name = ("c", 456)

def model(d2_name: Name = EXAMPLE_D2_NAME_1) -> None:
    # Note how our assertion is no longer needed, 
    # because we hardcoded the data in our desired format to begin with.
    d2 = d2_name[0]
    print(extract_literal(d2))

model(EXAMPLE_D2_NAME_2)

如果不是这样，那听起来你是从外部获取这个值。在这种情况下，
“解析，而不是验证”这个概念就很重要。为了更好地利用类型系统和mypy，我们应该把数据解析成我们期望的格式，然后再检查是否符合条件。可以看看mypy playground。

def parse_name(raw: str) -> Name:
    components = raw.split("-")

    if len(components) != 2:
        raise AssertionError("Exactly one hyphen expected.")

    prefix, raw_number = components

    try:
        number = int(components[1])
    except ValueError:   
        raise AssertionError("Expected number after hyphen.")
      
    if prefix in ("b", "c"):
        # Since `mypy` doesn't support literal narrowing yet: https://github.com/python/mypy/issues/16820,
        return prefix, number # type: ignore

    raise AssertionError("Prefix must be b or c.")

external_d2_name = parse_name(os.environ["D2_NAME"])
model(external_d2_name)

注意到我解析数据的方法和你验证数据的方法都使用了AssertionError，但我还在顺利的情况下返回了经过类型检查的数据，这样可以向其他代码证明它是正确的格式。

在评论中，有人提到TypeGuard。我不会使用这个，因为mypy实际上并不会检查你在TypeGuard中返回的内容。这个机制更适合用于可变长度的容器（比如list）。