如何实现接收Literal类型并返回该类型的断言函数？

想要编写一个断言函数，验证用户输入是否匹配预定义的Literal类型。例如：

MyLiteral = Literal["foo", "bar"]

期望通过以下方式调用函数：

some_user_provided_value = input()   # 示例

good_value = assert_literal(MyLiteral, some_user_provided_value)

要求：

• good_value的类型能被推断为MyLiteral
• 若输入不匹配Literal定义的任何值，抛出AssertionError

该函数等价于以下逻辑：

some_user_provided_value = input()   # 示例

good_value: MyLiteral

if some_user_provided_value == "foo":
    good_value = "foo"
elif some_user_provided_value == "bar":
    good_value = "bar"
else:
    raise AssertionError(f"Value {some_user_provided_value!r} is not a MyLiteral")

尝试的实现代码无法正常工作：

from typing import Any, TypeVar

T = TypeVar("T")


def assert_literal(literal_type: T, value: Any) -> T:
    if value not in typing.get_args(literal_type):
        raise AssertionError(f"Value {value!r} is not a {literal_type!r}")

    return typing.cast(T, value)

问题在于：该函数期望接收T的实例并返回T的实例，但MyLiteral是类型而非实例；若改用Type[T]，又因为"foo"并非MyLiteral的实例而失效。

询问是否有办法实现需求，以及Literal类型能否与TypeVar结合使用。

可以通过正确的类型标注结合typing.get_args来实现这个功能，具体步骤如下：

1. 导入必要模块：需要TypeVar、get_args等工具，以及Literal本身。
2. 定义绑定类型的TypeVar：让TypeVar关联到Literal类型，确保类型检查器能正确推断返回类型。
3. 实现断言逻辑：通过get_args获取Literal定义的所有有效值，验证输入是否在其中，通过则返回值（类型检查器会自动推断为目标Literal类型）。

完整实现代码：

from typing import Any, Literal, TypeVar, get_args

# 定义TypeVar，用于关联Literal类型
T = TypeVar("T")

def assert_literal(literal_type: type[T], value: Any) -> T:
    # 获取Literal类型定义的所有有效值
    valid_values = get_args(literal_type)
    
    # 校验是否为有效的Literal类型（至少有一个值）
    if not valid_values:
        raise TypeError("参数必须是包含至少一个值的Literal类型")
    
    # 验证输入值是否在允许范围内
    if value not in valid_values:
        raise AssertionError(
            f"值 {value!r} 不是有效的 {literal_type.__name__}，允许的值为 {valid_values}"
        )
    
    # 返回值，类型检查器会自动推断为对应的Literal类型
    return value

使用示例 #

MyLiteral = Literal["foo", "bar"]
some_user_provided_value = input()

# 调用函数后，good_value的类型会被推断为MyLiteral
good_value = assert_literal(MyLiteral, some_user_provided_value)

关键说明 #

• 类型标注：用type[T]表示接收的是T的类型对象（这里T就是我们定义的Literal类型），返回类型设为T，让类型检查器能正确推断返回值的类型。
• get_args的作用：该函数能提取出Literal类型中定义的所有有效值（比如MyLiteral对应的("foo", "bar")），以此作为验证依据。
• 兼容性：该实现支持Python 3.8及以上版本（Literal从3.8开始引入），且能被主流类型检查器（如mypy、pyright）正确识别。

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