Conditionally set FastAPI response model for route

Question

I'm trying to return a list of objects of type Company, including only "approved" ones, and with more or less attributes depending on whether the user requesting the list is a superuser or a regular user. This is my code so far:

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

The function correctly returns the objects according to request (ie., only is_approved=True companies if a regular request, and both is_approved=True and is_approved=False if requested by a superuser. Problem is, both cases use schema.CompanyRegularUsers, and I'd like to use schema.CompanySuperusers when SU's make the request.

How can I achieve that feature? I.e, is there a way to conditionally set the response_model property of the decorator function?

I've tried using JSONResponse and calling Pydantic's schema.CompanySuperusers.from_orm(), but it won't work with a list of Companies...

Answer 1

You can try to use Union type operator.

Your code would become

from typing import Union

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

this way, you specify as response model a list of either schema.CompanyRegularUsers or schema.CompanySuperUser

Let me know if it works, since I didn't test it