如何为支持A/B/M任意组合的price_type参数添加类型提示？

静态类型工具（比如mypy、pyright）没法直接描述“仅允许A/B/M任意组合”这类规则，所以不用费劲枚举所有可能的字符串组合，推荐用「语义化类型标记 + 运行时校验」的组合方案，既保留类型提示的可读性，又能严格约束输入。以下是几个实用实现：

方案1：原生NewType+运行时校验 #

用NewType给字符串赋予特定语义，配合简单的运行时检查确保输入合法，完全依赖Python标准库，无需额外依赖：

from typing import NewType

# 定义语义化类型，明确这是price_type参数的专属类型
PriceType = NewType("PriceType", str)

def __get_base_data(
    symbol: str,
    count: int,
    price_type: PriceType, 
):
    # 运行时校验：检查所有字符都是A/B/M，且无重复（若允许重复可去掉去重逻辑）
    valid_chars = {'A', 'B', 'M'}
    input_chars = set(price_type)
    if not input_chars.issubset(valid_chars):
        invalid_chars = input_chars - valid_chars
        raise ValueError(f"price_type包含无效字符：{''.join(invalid_chars)}，仅允许A/B/M的任意组合")
    
    # 函数核心逻辑
    pass

# 使用时可以显式转换（增强类型提示的严谨性，也可以直接传符合规则的字符串）
__get_base_data("BTCUSDT", 100, PriceType("AB"))
__get_base_data("BTCUSDT", 100, PriceType("BA"))  # 合法，顺序不影响

方案2：Annotated+自定义校验（Python 3.9+） #

Python 3.9及以上可以用Annotated把类型和校验逻辑绑定，静态检查工具（如pyright）能识别类型语义，运行时也能自动执行校验：

from typing import Annotated

def validate_price_type(value: str) -> str:
    valid_chars = {'A', 'B', 'M'}
    input_chars = set(value)
    if not input_chars.issubset(valid_chars):
        invalid_chars = input_chars - valid_chars
        raise ValueError(f"price_type无效：仅允许A/B/M的任意组合，当前包含非法字符{''.join(invalid_chars)}")
    return value

# 用Annotated标记类型和对应的校验函数
PriceType = Annotated[str, validate_price_type]

def __get_base_data(
    symbol: str,
    count: int,
    price_type: PriceType, 
):
    pass

方案3：借助Pydantic（适合已有Pydantic的项目） #

如果项目已经在用Pydantic做数据校验，可以用它的字符串约束功能，一行代码搞定类型提示和自动校验：

from pydantic import constr, ValidationError

# 用正则约束：仅允许A/B/M组成的任意长度字符串（若要禁止重复字符，可把regex改成r'^(?:([ABM])(?!.*\1))+$'）
PriceType = constr(regex=r'^[ABM]+$')

def __get_base_data(
    symbol: str,
    count: int,
    price_type: PriceType, 
):
    pass

# 测试示例
try:
    __get_base_data("BTCUSDT", 100, "ABM")  # 合法
    __get_base_data("BTCUSDT", 100, "BA")   # 合法
    __get_base_data("BTCUSDT", 100, "AX")   # 触发ValidationError
except ValidationError as e:
    print(e)

关键提醒 #

静态类型系统本身没法完全表达“任意A/B/M组合”的规则，所以无论用哪种方案，运行时校验都是必须的——它能在函数执行前拦截非法输入，避免后续逻辑出问题。

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