custom_api_response_type.py

"""use pydantic version 2"""

import enum
from typing import Generic, TypeVar

from pydantic import Field, BaseModel
from fastapi.encoders import jsonable_encoder


 T = TypeVar("T")


class ApiStatus(enum):
  SUCCESS = 1
  FAILURE = 0

  
class PaginatedContent(BaseModel, Generic[T]):
    """ Content data type for lists with pagination"""
    data: T
    total_count: int = 0
    limit: int = 100
    offset: int = 0
  
  
class ApiResponseHeader(BaseModel, Generic[T]):
    """ Header type of APIResponseType"""
    status: int = Field(..., description=str(list(ApiStatus)))
    message: str = "Success"
      
      
class APIResponseType(BaseModel, Generic[T]):
    """
    an api response type for using as the api's router response_model
    """
    header: ApiResponseHeader
    content: T

class APIResponse(Generic[T]):
      """
      Custom reponse class for apis
      Adds custom header, messages to reponses
      """

      header: ApiResponseHeader
      content: dict | list[dict] | None = None

      def __init__(
          self,
          data: T,
          *args,
          msg_code: int = 0,
          msg_status: ApiStatus = ApiStatus.SUCCESS,
          **kwargs
      ) -> None:
          self.header = ApiResponseHeader(
              status=msg_status,
              message="",
          )
          if isinstance(data, BaseModel):
              self.content = data.model_dump()
          else:
              self.content = jsonable_encoder(data)