如何为Pydantic模型的参数设置特定约束？

针对你提出的参数约束需求，以下是修改后的代码实现，同时修正原代码中code和cluster的类型错误（原代码用Set[str]表示集合，但需求是单个值属于指定集合）：

核心修改说明 #

• 日期字段改用date类型，Pydantic自动解析字符串日期，便于日期比较
• 通过Pydantic的validator实现日期约束逻辑
• 用Literal或Enum限定code和cluster的可选值，确保输入值只能是指定集合中的单个值

方式一：使用Literal直接限定可选值 #

from fastapi import FastAPI
from pydantic import BaseModel, Field, validator
from datetime import date
from typing import Literal
import uvicorn

app = FastAPI()

class Query(BaseModel):
    start_date: date = Field(..., description="起始日期，不能早于2019-01-01")
    end_date: date = Field(..., description="结束日期，必须晚于起始日期")
    code: Literal[
        "A1", "A2", "A3", "A4",
        "X1", "X2", "X3", "X4", "X5",
        "Y1", "Y2", "Y3"
    ] = Field(..., description="只能是指定集合中的单个值")
    cluster: Literal["C1", "C2", "C3"] = Field(..., description="只能是指定集合中的单个值")

    @validator('start_date')
    def validate_start_date(cls, value):
        min_date = date(2019, 1, 1)
        if value < min_date:
            raise ValueError(f"起始日期不能早于{min_date.isoformat()}")
        return value

    @validator('end_date')
    def validate_end_date(cls, value, values):
        if 'start_date' in values and value <= values['start_date']:
            raise ValueError("结束日期必须晚于起始日期")
        return value

@app.post("/")
async def read_table(query: Query):
    return {"msg": query.dict()}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

方式二：使用Enum（推荐值较多或需复用场景） #

from fastapi import FastAPI
from pydantic import BaseModel, Field, validator
from datetime import date
from enum import Enum
import uvicorn

app = FastAPI()

class CodeEnum(str, Enum):
    A1 = "A1"
    A2 = "A2"
    A3 = "A3"
    A4 = "A4"
    X1 = "X1"
    X2 = "X2"
    X3 = "X3"
    X4 = "X4"
    X5 = "X5"
    Y1 = "Y1"
    Y2 = "Y2"
    Y3 = "Y3"

class ClusterEnum(str, Enum):
    C1 = "C1"
    C2 = "C2"
    C3 = "C3"

class Query(BaseModel):
    start_date: date = Field(..., description="起始日期，不能早于2019-01-01")
    end_date: date = Field(..., description="结束日期，必须晚于起始日期")
    code: CodeEnum = Field(..., description="只能是指定集合中的单个值")
    cluster: ClusterEnum = Field(..., description="只能是指定集合中的单个值")

    @validator('start_date')
    def validate_start_date(cls, value):
        min_date = date(2019, 1, 1)
        if value < min_date:
            raise ValueError(f"起始日期不能早于{min_date.isoformat()}")
        return value

    @validator('end_date')
    def validate_end_date(cls, value, values):
        if 'start_date' in values and value <= values['start_date']:
            raise ValueError("结束日期必须晚于起始日期")
        return value

@app.post("/")
async def read_table(query: Query):
    return {"msg": query.dict()}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

关键约束说明 #

1. start_date约束：通过validator检查输入日期是否早于2019-01-01，不符合则抛出错误
2. end_date约束：依赖start_date的值，确保结束日期严格晚于起始日期
3. code/cluster约束：通过Literal或Enum限定输入值只能是指定集合中的单个值，Pydantic会自动验证输入合法性

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