如何正确实现 pydantic dict 键作为 Enum/Literal 以支持 swagger

Question

我正在尝试编写带有预定义键的字典的 pydantic 模型：

class Result(BaseModel):
    name: str = Field(title="Name")
    description: str = Field(title="Description")
    status: Literal["WARN", "PASS", "FAIL"] = Field(title="Status")


class Response(BaseModel):
    results: dict[Literal["WARN", "PASS", "FAIL"], List[Result]] = Field(title="title")


@app.get("/test", response_model=Response)
def test():

    # this is recevided from external package
    res = {'PASS': [{'name': 'name', 'description': 'desc', 'status': 'PASS'}], 
           'WARN': [],
           'FAIL': []
           }

    response = Response(
        results=res
    )
    return response

但在 Swagger 文档中它显示为

< * >

:

有趣的是，它可以显示文字（例如

status

），但不能显示为字典键。

Answer 1

为什么不简单地定义显式属性呢？像这样：

from pydantic import BaseModel
from enum import StrEnum


class Status(StrEnum):
    WARN = "WARN"
    PASS = "PASS"
    FAIL = "FAIL"


class Result(BaseModel):
    name: str
    description: str
    status: Status


class Results(BaseModel):
    WARN: list[Result] = []
    PASS: list[Result] = []
    FAIL: list[Result] = []


class Response(BaseModel):
    results: Results

这将使您在序列化/反序列化方面获得与基于

dict

的解决方案相同的行为：

>>> r = Response.model_validate_json('''{"results": {"WARN": [{"name": "example", "description": "example", "status": "WARN"}]}}''')
>>> r
Response(results=Results(WARN=[Result(name='example', description='example', status=<Status.WARN: 'WARN'>)], PASS=[], FAIL=[]))
>>> print(r.model_dump_json(indent=2))
{
  "results": {
    "WARN": [
      {
        "name": "example",
        "description": "example",
        "status": "WARN"
      }
    ],
    "PASS": [],
    "FAIL": []
  }
}

它会产生一个 JSONSchema，它限制

results

键上可用的属性：

{
  "$defs": {
    .
    .
    .
    "Results": {
      "properties": {
        "WARN": {
          "default": [],
          "items": {
            "$ref": "#/$defs/Result"
          },
          "title": "Warn",
          "type": "array"
        },
        "PASS": {
          "default": [],
          "items": {
            "$ref": "#/$defs/Result"
          },
          "title": "Pass",
          "type": "array"
        },
        "FAIL": {
          "default": [],
          "items": {
            "$ref": "#/$defs/Result"
          },
          "title": "Fail",
          "type": "array"
        }
      },
      "title": "Results",
      "type": "object"
    },
  "properties": {
    "results": {
      "$ref": "#/$defs/Results"
    }
  },
  "required": [
    "results"
  ],
  "title": "Response",
  "type": "object"
}