前言

众所周知，Python 的流行Web后端服务框架 FastAPI 自带有在线文档，位于 /docs 和 /redoc，同时也提供 openapi.json，位于 /openapi.json。FastAPI 会根据开发者为每个路由编写的方法的有关信息来生成 openapi.json，比如路由名称取自方法名称、文档取自方法的 docstring。

这件事唯一的小问题点在于，FastAPI 不会解析 docstring，而是把 docstring 原封不动地写到 openapi.json 中。然后，文档基于 openapi.json 渲染，会导致令人很不爽的格式丢失。虽然 openapi.json 标准支持此处的文档使用 Markdown，但是至少芒果没有在 docstring 中使用 Markdown 美化排版的习惯，并且那样也会让 PyCharm 摸不着头脑。

然后在其他可以使用 openapi.json 生成文档的东西那里，这样生成出来的文档也很让人头大。

我们很多时候会希望 docstring 中的内容能在正确的地方显示，比如让 Args: 后的参数文档在 openapi.json 的参数部分显示。但是这需要 FastAPI 对 docstring 进行解析，而 FastAPI 在近日关闭了一个与此有关的 PR，他们的态度是不进行这项改动。

因此显然，如果想要 openapi.json 更加合理的生成，我们需要自行解决从 docstring 到 openapi.json 的这一步骤。

派遣 AI！

芒果现在比较倾向于 Google 风格的 docstring，即：

    """
    修改聊天室的基本信息。
    content 和 script 需要从各自的独立端点请求修改，不包含在本端点的负责范围内。

    Args:
        id_: 聊天室 ID
        chatroom:  聊天室基本信息，类型为 ChatroomPublic。注意 id 不可更改，如提供则会被忽略
        user: 用户
        session: 数据库连接对象

    Raises:
        HTTPException: 404 表示未找到聊天室

    Returns:
        修改过的聊天室对象，供前端更新。
    """

然后压力 AI 帮我写。

需求是将 Google 风格的 docstring 中的 Args、Raises、Returns 等块放到他们该在的地方去，让 kimi 先写，然后让 Deepseek review 一遍，目前测试没有问题。

奇怪的是，毕竟是用来解析 Google 风格 docstring 的东西，AI 自己使用的却还是 Numpy 风格的表格文档，看来多少是有点逆反在里面的呢喵。

"""
FastAPI Docstring-to-OpenAPI Enricher
=====================================

自动解析 Google Style docstring，将 Args / Returns / Raises 注入到 OpenAPI Schema 的对应位置，
让 vitepress-openapi 等工具能够正确渲染参数表格和返回值说明。

功能特性
--------
1. 参数描述注入：路径参数、查询参数、Header、Cookie、请求体字段
2. 返回值描述注入：自动写入 2xx 响应的 description 及 schema.description
3. 异常描述注入：以 Markdown 格式附加到 operation description
4. 嵌套 Schema 递归处理：支持 Pydantic v1/v2 的嵌套模型、List、Optional 等
5. \f 截断兼容：与 FastAPI 原生行为保持一致
6. 全局 Schema 补全：为 components/schemas 中缺少描述的字段补充 docstring 说明
7. 零侵入：只需替换 ``app.openapi``，业务代码无需任何修改

依赖安装
--------
    pip install docstring-parser

使用方式
--------
    from fastapi import FastAPI
    from fastapi_docstring_openapi import enrich_openapi_from_docstrings

    app = FastAPI()
    enrich_openapi_from_docstrings(app)

    # 你的路由注册...

完整示例见文件底部 ``if __name__ == "__main__":`` 块。
"""

from __future__ import annotations

import inspect
import logging
from typing import Any, Dict, Optional, Set

from docstring_parser import DocstringStyle, ParseError, parse
from fastapi import FastAPI
from fastapi.openapi.utils import get_openapi

logger = logging.getLogger(__name__)


def _parse_docstring(
    func: Any, style: DocstringStyle = DocstringStyle.AUTO
) -> Optional[Any]:
    """解析函数的 Google Style docstring，失败时返回 None。"""
    doc = inspect.getdoc(func)
    if not doc:
        return None
    if "\f" in doc:
        doc = doc.split("\f")[0].strip()
    try:
        return parse(doc, style=style)
    except ParseError:
        return None
    except Exception:
        logger.warning(
            "Unexpected error parsing docstring for %s",
            getattr(func, "__name__", func),
            exc_info=True,
        )
        return None


def _build_param_lookup(parsed_doc: Any) -> Dict[str, str]:
    """从解析后的 docstring 构建 {参数名: 描述} 映射。"""
    lookup: Dict[str, str] = {}
    for param in getattr(parsed_doc, "params", []) or []:
        name = getattr(param, "arg_name", None)
        desc = getattr(param, "description", None) or ""
        if name:
            # 合并多行描述并清理缩进
            lookup[name] = _dedent_description(desc)
    return lookup


def _build_returns_description(parsed_doc: Any) -> Optional[str]:
    """从解析后的 docstring 构建返回值描述。"""
    ret = getattr(parsed_doc, "returns", None)
    if not ret:
        return None
    type_name = getattr(ret, "type_name", None) or ""
    desc = getattr(ret, "description", None) or ""
    if type_name and desc:
        return f"**{type_name}**: {desc}"
    return desc or None


def _build_raises_description(parsed_doc: Any) -> Optional[str]:
    """从解析后的 docstring 构建异常描述（Markdown 列表）。"""
    raises = getattr(parsed_doc, "raises", None)
    if not raises:
        return None
    parts: list[str] = []
    for exc in raises:
        type_name = getattr(exc, "type_name", None) or ""
        desc = getattr(exc, "description", None) or ""
        if type_name and desc:
            parts.append(f"- **{type_name}**: {desc}")
        elif type_name:
            parts.append(f"- **{type_name}**")
        elif desc:
            parts.append(f"- {desc}")
    return "\n".join(parts) if parts else None


def _dedent_description(text: str) -> str:
    """清理 docstring 描述中的多余缩进，保留段落结构。"""
    if not text:
        return ""
    lines = text.splitlines()
    # 找到最小公共缩进（排除空行）
    min_indent = min(
        (len(line) - len(line.lstrip()) for line in lines if line.strip()),
        default=0,
    )
    cleaned = [line[min_indent:] if line.strip() else "" for line in lines]
    return "\n".join(cleaned).strip()


def _get_clean_operation_description(parsed_doc: Any) -> str:
    """
    生成干净的 operation description。

    策略：保留 summary + long_description，去掉 Args/Returns/Raises 等机器块。
    """
    parts: list[str] = []
    short = getattr(parsed_doc, "short_description", None)
    long_ = getattr(parsed_doc, "long_description", None)

    if short:
        parts.append(short)
    if long_:
        parts.append(long_)

    return "\n\n".join(parts).strip()


def _inject_schema_properties(
        schema: Dict[str, Any],
        param_lookup: Dict[str, str],
        visited_refs: Optional[Set[str]] = None,
        openapi_schema: Optional[Dict[str, Any]] = None,
) -> None:
    """
    递归注入 schema properties 的描述。

    支持对象、数组、allOf/anyOf/oneOf、$ref 引用（全局 components/schemas 补全）。
    """
    if visited_refs is None:
        visited_refs = set()
    if not isinstance(schema, dict):
        return

    # 处理 $ref：在 components/schemas 中查找并补全
    ref = schema.get("$ref")
    if ref and openapi_schema:
        if ref not in visited_refs:
            visited_refs.add(ref)
            # 提取 schema 名，例如 "#/components/schemas/User" -> "User"
            schema_name = ref.split("/")[-1]
            components = openapi_schema.get("components", {}).get("schemas", {})
            target = components.get(schema_name)
            if target:
                _inject_schema_properties(target, param_lookup, visited_refs, openapi_schema)
        return

    # 处理 properties（对象字段）
    properties = schema.get("properties")
    if isinstance(properties, dict):
        for prop_name, prop_schema in properties.items():
            if prop_name in param_lookup:
                # 只有当 docstring 描述非空时才写入
                if param_lookup[prop_name]:
                    prop_schema["description"] = param_lookup[prop_name]
            # 递归处理子 schema
            _inject_schema_properties(prop_schema, param_lookup, visited_refs, openapi_schema)

    # 处理 items（数组类型）
    items = schema.get("items")
    if items:
        _inject_schema_properties(items, param_lookup, visited_refs, openapi_schema)

    # 处理 allOf / anyOf / oneOf
    for key in ("allOf", "anyOf", "oneOf"):
        for sub in schema.get(key, []):
            _inject_schema_properties(sub, param_lookup, visited_refs, openapi_schema)

    # 处理 additionalProperties
    additional = schema.get("additionalProperties")
    if isinstance(additional, dict):
        _inject_schema_properties(additional, param_lookup, visited_refs, openapi_schema)


def enrich_openapi_from_docstrings(
        app: FastAPI,
        *,
        prefer_docstring_over_field: bool = True,
        append_raises_to_description: bool = True,
        docstring_style: DocstringStyle = DocstringStyle.AUTO,
) -> None:
    """
    为 FastAPI 应用启用 docstring 驱动的 OpenAPI 增强。

    参数
    ----
    app : FastAPI
        目标应用实例。
    prefer_docstring_over_field : bool, 默认 True
        当 docstring 中的参数描述与 Pydantic Field(description=...) 冲突时，
        是否优先使用 docstring 的描述。
    append_raises_to_description : bool, 默认 True
        是否将 Raises 段落以 Markdown 格式追加到 operation description。
    docstring_style : DocstringStyle, 默认 AUTO
        解析风格。团队统一用 Google 时可显式传入 ``DocstringStyle.GOOGLE``。
    """

    # 保存原始的 openapi 函数（如果有）
    original_openapi = app.openapi

    def custom_openapi() -> Dict[str, Any]:
        # 如果已经有缓存，直接返回
        if app.openapi_schema:
            return app.openapi_schema

        # 生成基础 schema
        openapi_schema = get_openapi(
            title=app.title,
            version=app.version,
            openapi_version=app.openapi_version,
            description=app.description,
            routes=app.routes,
        )

        for route in app.routes:
            if not hasattr(route, "endpoint") or not route.endpoint:
                continue

            path = getattr(route, "path", None)
            methods = [m.lower() for m in route.methods] if hasattr(route, "methods") else []
            if not path or not methods:
                continue

            # 解析 docstring
            parsed = _parse_docstring(route.endpoint, style=docstring_style)
            if not parsed:
                continue

            param_lookup = _build_param_lookup(parsed)
            returns_desc = _build_returns_description(parsed)
            raises_desc = _build_raises_description(parsed)
            clean_desc = _get_clean_operation_description(parsed)

            for method in methods:
                if method == "head":
                    # FastAPI 通常不把 HEAD 单独写入 OpenAPI，或者与 GET 共享
                    continue
                if method not in openapi_schema.get("paths", {}).get(path, {}):
                    continue

                op = openapi_schema["paths"][path][method]

                # ---------- 1. 更新 operation description ----------
                if clean_desc:
                    op["description"] = clean_desc

                # ---------- 2. 注入参数描述（路径/查询/Header/Cookie）----------
                for param in op.get("parameters", []):
                    param_name = param.get("name")
                    if param_name and param_name in param_lookup:
                        existing = param.get("description", "")
                        new_desc = param_lookup[param_name]
                        if new_desc and (prefer_docstring_over_field or not existing):
                            param["description"] = new_desc

                # ---------- 3. 注入请求体 schema 字段描述 ----------
                if "requestBody" in op:
                    content = op["requestBody"].get("content", {})
                    for media_obj in content.values():
                        schema = media_obj.get("schema", {})
                        if schema:
                            _inject_schema_properties(
                                schema, param_lookup, openapi_schema=openapi_schema
                            )

                # ---------- 4. 注入返回值描述 ----------
                if returns_desc:
                    for code, resp in op.get("responses", {}).items():
                        if code.startswith("2"):  # 2xx 成功响应
                            # 更新响应级 description
                            resp["description"] = returns_desc
                            # 同时注入到响应 schema 的 description（如果存在）
                            for media_obj in resp.get("content", {}).values():
                                schema = media_obj.get("schema", {})
                                if schema and not schema.get("description"):
                                    schema["description"] = returns_desc
                                    # 递归注入 schema 内部字段
                                    _inject_schema_properties(
                                        schema, param_lookup, openapi_schema=openapi_schema
                                    )

                # ---------- 5. 异常描述追加 ----------
                if append_raises_to_description and raises_desc:
                    current_desc = op.get("description", "")
                    raises_md = f"## 异常\n\n{raises_desc}"
                    # 避免重复追加
                    if raises_md not in current_desc:
                        op["description"] = f"{current_desc}\n\n{raises_md}".strip()

        app.openapi_schema = openapi_schema
        return app.openapi_schema

    app.openapi = custom_openapi

然后，对于开发者，只需要在 FastAPI 对象建立后多增加一行即可：

# 当然得先导入这个方法，看你怎么规划你的代码咯
enrich_openapi_from_docstrings(app)

大功告成！~