如何校验字符串，让mypy识别其为指定Literal类型？

方法1：类型断言（简单直接） #

在完成合法性校验后，直接通过断言或类型转换告诉mypy此时market的类型为Markets——因为我们已经确保它属于合法取值范围，所以操作是安全的。

用assert实现 #

import os
from typing import Literal, get_args

Markets = Literal[
    "BE", "DE", "DK", "EE", "ES", "FI", "FR", "GB", "IT", "LT", "LV", "NL", "NO", "PL", "PT", "SE"
]
MARKETS: list[Markets] = list(get_args(Markets))


def foo(x: Markets) -> None:
    print(x)


market = os.environ.get("market")

if market not in MARKETS:
    raise ValueError("Invalid market")

# 用assert让mypy识别类型
assert market in MARKETS
foo(market)

用cast实现 #

如果担心assert在生产环境被优化掉，可以用cast强制转换类型：

import os
from typing import Literal, get_args, cast

Markets = Literal[
    "BE", "DE", "DK", "EE", "ES", "FI", "FR", "GB", "IT", "LT", "LV", "NL", "NO", "PL", "PT", "SE"
]
MARKETS: list[Markets] = list(get_args(Markets))


def foo(x: Markets) -> None:
    print(x)


market = os.environ.get("market")

if market not in MARKETS:
    raise ValueError("Invalid market")

# 强制转换类型
market = cast(Markets, market)
foo(market)

方法2：类型守卫函数（更严谨，适合复用） #

如果需要在多处做相同的市场校验，写一个类型守卫函数，让mypy能自动推导变量类型：

import os
from typing import Literal, get_args, TypeGuard

Markets = Literal[
    "BE", "DE", "DK", "EE", "ES", "FI", "FR", "GB", "IT", "LT", "LV", "NL", "NO", "PL", "PT", "SE"
]
MARKETS: list[Markets] = list(get_args(Markets))

# 类型守卫函数，返回True时mypy会自动将参数窄化为Markets类型
def is_valid_market(market: str | None) -> TypeGuard[Markets]:
    return market in MARKETS

def foo(x: Markets) -> None:
    print(x)


market = os.environ.get("market")

if not is_valid_market(market):
    raise ValueError("Invalid market")

# 此时mypy会自动识别market为Markets类型
foo(market)

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