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

Question

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

@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)
#

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

我怎样才能实现这个功能？即，有没有办法有条件地设置response_model装饰器函数的属性？

我尝试过使用JSONResponse和调用 Pydantic's 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_encoderFastAPI 在每个默认响应的底层使用的编码器来序列化列表。

编辑：错别字