将字符串转为Literal类型时mypy报类型不兼容错误，如何修复？

问题场景 #

运行mypy 1.9.0检查以下Python 3.11代码时，出现类型不兼容错误：

代码片段 #

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" 格式，需根据首字母输出对应信息。

解决方案 #

mypy无法自动识别if d2_name[0] not in ["b", "c"]: raise逻辑对类型的约束，因此需要手动帮助mypy确认d2的类型，以下是三种可行方法：

方法1：显式标注Literal类型 #

直接给d2标注Literal["b", "c"]类型，明确告诉mypy变量的取值范围：

def model(d2_name: str = "b-123") -> None:
    if d2_name[0] not in ["b", "c"]:
        raise AssertionError
    d2: Final[Literal["b", "c"]] = d2_name[0]  # 显式标注类型
    print(extract_literal(d2))

方法2：使用cast函数做类型断言 #

通过cast函数强制指定变量类型，适合需要明确告知类型检查器的场景：

from typing import Literal, Final, cast  # 导入cast

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 = cast(Literal["b", "c"], d2_name[0])  # 类型断言
    print(extract_literal(d2))

方法3：用assert语句优化类型推断 #

将原有的if raise逻辑替换为assert，mypy会自动识别断言后的类型约束，无需额外标注：

def model(d2_name: str = "b-123") -> None:
    assert d2_name[0] in ("b", "c"), "d2_name必须以'b'或'c'开头"
    d2: Final = d2_name[0]
    print(extract_literal(d2))

原因说明 #

mypy的类型推断无法从if ... raise语句中自动缩小变量类型，但能识别assert语句的约束；显式标注类型或使用cast则是直接绕过类型检查器的自动推断，手动指定类型，这两种方式都能解决类型不兼容的问题。

内容的提问来源于stack exchange，提问作者ashnair1