仅当作为 FastAPI 调用的一部分返回时才排除 pydantic 模型字段

rbh*_*lla 4 python starlette pydantic fastapi

语境

我有一个非常复杂的 pydantic 模型,其中有很多嵌套的 pydantic 模型。我想确保某些字段永远不会作为 API 调用的一部分返回,但我希望这些字段用于内部逻辑。

我尝试过的

我首先尝试使用 pydantic 的Field函数来指定exclude我不想返回的字段上的标志。这有效,但是我的内部逻辑中的函数每当.dict()通过调用.dict(exclude=None).

相反,我return_in_api在 上指定了一个自定义标志Field,目标是仅在 FastAPI 调用时应用排除.dict()。我尝试编写一个中间件来根据包含的嵌套字段来调用.dict()和传递我自己的属性。然而,FastAPI 的中间件为我提供了一个响应流,我不想过早地解析它。excludereturn_in_api=False

相反,我编写了一个装饰器,它.dict()使用适当的值调用路由处理程序的返回值exclude

问题

一个挑战是,每当添加新端点时,添加它们的人都必须记住包含此装饰器,否则字段会泄漏。

理想情况下,我想将此装饰器应用于每个路由,但通过中间件执行此操作似乎会破坏响应流。

Dan*_*erg 6

系统地排除所有路线的字段

我发现最好使用一个具体但非常简单的示例。假设您有以下模型:

from pydantic import BaseModel, Field


class Demo(BaseModel):
    foo: str
    bar: str = Field(return_in_api=False)
Run Code Online (Sandbox Code Playgroud)

我们希望确保bar永远不会在响应中返回,无论是response_model显式提供为路由装饰器的参数还是仅将其设置为路由处理函数的返回注释时。exclude(假设我们出于某种原因不想在字段中使用内置参数。)

我发现最可靠的方法是子类化fastapi.routing.APIRoute并挂钩到它的__init__方法。通过复制父类的一小部分代码,我们可以确保我们始终获得正确的响应模型。一旦我们有了这个,只需在response_model_exclude调用父构造函数之前设置路由的参数即可。

这是我的建议:

from collections.abc import Callable
from typing import Any

from fastapi.responses import Response
from fastapi.dependencies.utils import get_typed_return_annotation, lenient_issubclass
from fastapi.routing import APIRoute, Default, DefaultPlaceholder


class CustomAPIRoute(APIRoute):
    def __init__(
        self,
        path: str,
        endpoint: Callable[..., Any],
        *,
        response_model: Any = Default(None),
        **kwargs: Any,
    ) -> None:
        # We need this part to ensure we get the response model,
        # even if it is just set as an annotation on the handler function.
        if isinstance(response_model, DefaultPlaceholder):
            return_annotation = get_typed_return_annotation(endpoint)
            if lenient_issubclass(return_annotation, Response):
                response_model = None
            else:
                response_model = return_annotation
        # Find the fields to exclude:
        if response_model is not None:
            kwargs["response_model_exclude"] = {
                name
                for name, field in response_model.__fields__.items()
                if field.field_info.extra.get("return_in_api") is False
            }
        super().__init__(path, endpoint, response_model=response_model, **kwargs)
Run Code Online (Sandbox Code Playgroud)

我们现在可以在路由器上设置自定义路由类(文档)。这样它将用于其所有路线:

from fastapi import FastAPI
# ... import CustomAPIRoute
# ... import Demo

api = FastAPI()
api.router.route_class = CustomAPIRoute


@api.get("/demo1")
async def demo1() -> Demo:
    return Demo(foo="a", bar="b")


@api.get("/demo2", response_model=Demo)
async def demo2() -> dict[str, Any]:
    return {"foo": "x", "bar": "y"}
Run Code Online (Sandbox Code Playgroud)

尝试使用uvicornGET设置端点这个简单的 API 示例/demo1,并分别/demo2产生响应{"foo":"a"}{"foo":"x"}


确保架构一致性

然而值得一提的是(除非我们采取额外的步骤)该bar字段仍然是 schema 的一部分。这意味着,例如,为这两个端点自动生成的 OpenAPI 文档将显示 bar为预期响应的顶级属性。

这不是您问题的一部分,因此我认为您已经意识到这一点并正在采取措施确保一致性。如果没有,对于阅读本文的其他人,您可以在基本模型schema_extra上定义一个静态方法Config,以删除那些在返回模式之前永远不会“向外部”显示的字段:

from typing import Any
from pydantic import BaseModel, Field


class CustomBaseModel(BaseModel):
    class Config:
        @staticmethod
        def schema_extra(schema: dict[str, Any]) -> None:
            properties = schema.get("properties", {})
            to_delete = set()
            for name, prop in properties.items():
                if prop.get("return_in_api") is False:
                    to_delete.add(name)
            for name in to_delete:
                del properties[name]


class Demo(CustomBaseModel):
    foo: str
    bar: str = Field(return_in_api=False)
Run Code Online (Sandbox Code Playgroud)