Fixing issues awtkns#153 & awtkns#138

Fa2y · Feb 22, 2023 · d7dea82 · d7dea82
1 parent 9b82986
commit d7dea82
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 7 deletions.
diff --git a/fastapi_crudrouter/core/_base.py b/fastapi_crudrouter/core/_base.py
@@ -4,7 +4,7 @@
 from fastapi import APIRouter, HTTPException
 from fastapi.types import DecoratedCallable
 
-from ._types import T, DEPENDENCIES
+from ._types import T, DEPENDENCIES, PAGINATIONEXTRADATA
 from ._utils import pagination_factory, schema_factory
 
 NOT_FOUND = HTTPException(404, "Item not found")
@@ -58,7 +58,7 @@ def __init__(
                 "",
                 self._get_all(),
                 methods=["GET"],
-                response_model=Optional[List[self.schema]],  # type: ignore
+                response_model=Optional[List[self.schema] | PAGINATIONEXTRADATA],  # type: ignore
                 summary="Get All",
                 dependencies=get_all_route,
             )

diff --git a/fastapi_crudrouter/core/_types.py b/fastapi_crudrouter/core/_types.py
@@ -1,4 +1,4 @@
-from typing import Dict, TypeVar, Optional, Sequence
+from typing import Dict, TypeVar, Optional, Sequence, List
 
 from fastapi.params import Depends
 from pydantic import BaseModel
@@ -8,3 +8,7 @@
 
 T = TypeVar("T", bound=BaseModel)
 DEPENDENCIES = Optional[Sequence[Depends]]
+
+class PAGINATIONEXTRADATA(BaseModel):
+    count: int
+    results: List
diff --git a/fastapi_crudrouter/core/tortoise.py b/fastapi_crudrouter/core/tortoise.py
@@ -1,7 +1,7 @@
 from typing import Any, Callable, List, Type, cast, Coroutine, Optional, Union
 
 from . import CRUDGenerator, NOT_FOUND
-from ._types import DEPENDENCIES, PAGINATION, PYDANTIC_SCHEMA as SCHEMA
+from ._types import PAGINATIONEXTRADATA, DEPENDENCIES, PAGINATION, PYDANTIC_SCHEMA as SCHEMA
 
 try:
     from tortoise.models import Model
@@ -13,7 +13,7 @@
 
 
 CALLABLE = Callable[..., Coroutine[Any, Any, Model]]
-CALLABLE_LIST = Callable[..., Coroutine[Any, Any, List[Model]]]
+CALLABLE_LIST = Callable[..., Coroutine[Any, Any, PAGINATIONEXTRADATA | List[Model]]]
 
 
 class TortoiseCRUDRouter(CRUDGenerator[SCHEMA]):
@@ -32,6 +32,7 @@ def __init__(
         update_route: Union[bool, DEPENDENCIES] = True,
         delete_one_route: Union[bool, DEPENDENCIES] = True,
         delete_all_route: Union[bool, DEPENDENCIES] = True,
+        paginationextradata: Union[bool, DEPENDENCIES] = False,
         **kwargs: Any
     ) -> None:
         assert (
@@ -58,19 +59,24 @@ def __init__(
         )
 
     def _get_all(self, *args: Any, **kwargs: Any) -> CALLABLE_LIST:
-        async def route(pagination: PAGINATION = self.pagination) -> List[Model]:
+        async def route(pagination: PAGINATION = self.pagination) -> PAGINATIONEXTRADATA | List[Model]:
             skip, limit = pagination.get("skip"), pagination.get("limit")
             query = self.db_model.all().offset(cast(int, skip))
+            if self.paginationextradata:
+                count = self.db_model.all().count() # added for issue #138
             if limit:
                 query = query.limit(limit)
+            query = self.schema.from_queryset(query) # added from issue #153
+            if self.paginationextradata:
+                return {"results": await query, "count": await count}
             return await query
 
         return route
 
     def _get_one(self, *args: Any, **kwargs: Any) -> CALLABLE:
         async def route(item_id: int) -> Model:
             model = await self.db_model.filter(id=item_id).first()
-
+            model = await self.schema.from_tortoise_orm(model) # added from issue #153
             if model:
                 return model
             else: