有条件地为路由设置 FastAPI 响应模型

Question

我正在尝试 return 公司类型的对象列表，仅包括“已批准”的对象，并且具有或多或少的属性，具体取决于请求该列表的用户是超级用户还是普通用户用户。到目前为止，这是我的代码：

@router.get("/", response_model=List[schema.CompanyRegularUsers])
def get_companies(db: Session = Depends(get_db), is_superuser: bool = Depends(check_is_superuser)):
    """
    If SU, also include sensitive data.
    """
    if is_superuser:
        return crud.get_companies_admin(db=db)
    return crud.get_companies_user(db=db)
#

该函数根据请求正确地 returns 对象（即，只有 is_approved=True 公司如果是常规请求，is_approved=True 和 is_approved=False 如果请求由超级用户。问题是，两种情况都使用 schema.CompanyRegularUsers，我想在 SU 发出请求时使用 schema.CompanySuperusers。

如何实现该功能？即，有没有办法有条件地设置装饰器函数的 response_model 属性？

我试过使用 JSONResponse 并调用 Pydantic 的 schema.CompanySuperusers.from_orm()，但它不适用于公司列表...

Answer 1

您可以尝试使用Union类型运算符。

您的代码将变成

from typing import Union

@router.get("/", response_model=List[Union[schema.CompanyRegularUsers, schema.CompanySuperUser]])

这样，您将 schema.CompanyRegularUsers 或 schema.CompanySuperUser

的列表指定为响应模型

让我知道它是否有效，因为我没有测试它

Answer 2

我最终通过返回自定义 JSONResponse 解决了这个谜题。它没有出现在自动文档中，但我想我可以在以后解决这个问题。代码如下，以防对别人有帮助：

...
from pydantic import parse_obj_as
from fastapi.responses import JSONResponse
from fastapi.encoders import jsonable_encoder
...

@router.get("/", response_model=List[schema.CompanyRegularUsers])
def get_companies(db: Session = Depends(get_db), is_superuser: bool = Depends(check_is_superuser)):
    """
    If SU, also include sensitive data.
    """
    if is_superuser:
        companies = parse_obj_as(List[schema.CompanyAdmin], crud.get_companies_admin(db=db))
        return JSONResponse(jsonable_encoder(companies))
    return crud.get_companies_user(db=db)

因此，在 is_admin 分支中，路径操作调用 pydantic 的 parse_obj_as 以便将 SQLAlchemy 的查询返回的对象列表映射为 CompanyAdmin 对象的列表。然后，它使用 jsonable_encoder，编码器 FastAPI 在每个默认响应中使用，来序列化列表。

编辑：打字错误

有条件地为路由设置 FastAPI 响应模型

Conditionally set FastAPI response model for route

python

orm

sqlalchemy

pydantic

fastapi