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

Question

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

@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 如果 super 用户请求。问题是，这两种情况都使用 schema.CompanyRegularUsers，我想在 SU 进行时使用 schema.CompanySuperusers请求。

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

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

Answer 1

我最终通过返回自定义 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 在内部为每个默认响应使用的编码器)来序列化列表。

编辑:打字错误