Added models where possible

classes/models.py

@@ -0,0 +1,21 @@
+from typing import Union
+from pydantic import BaseModel
+
+class Photo(BaseModel):
+    id: str
+    album: str
+    hash: str
+    filename: str
+
+class Album(BaseModel):
+    id: str
+    name: str
+    title: str
+
+class AlbumModified(BaseModel):
+    name: str
+    title: str
+
+class SearchResults(BaseModel):
+    results: list
+    next_page: Union[str, None] = None

extensions/albums.py

@@ -2,6 +2,7 @@ import re
 from os import makedirs, path, rename
 from shutil import rmtree
 from typing import Union
+from classes.models import Album, AlbumModified, SearchResults
 from modules.app import app
 from modules.database import col_photos, col_albums
 from modules.security import User, get_current_active_user
@@ -12,7 +13,7 @@ from fastapi import HTTPException, Security
 from fastapi.responses import UJSONResponse, Response
 from starlette.status import HTTP_204_NO_CONTENT, HTTP_404_NOT_FOUND, HTTP_406_NOT_ACCEPTABLE, HTTP_409_CONFLICT
 
-@app.post("/albums", response_class=UJSONResponse, description="Create album with name and title")
+@app.post("/albums", response_class=UJSONResponse, response_model=Album, description="Create album with name and title")
 async def album_create(name: str, title: str, current_user: User = Security(get_current_active_user, scopes=["albums.write"])):
 
     if re.search(re.compile('^[a-z,0-9,_]*$'), name) is False:
@@ -39,7 +40,7 @@ async def album_create(name: str, title: str, current_user: User = Security(get_
         }
     )
 
-@app.get("/albums", description="Find album by name")
+@app.get("/albums", response_model=SearchResults, description="Find album by name")
 async def album_find(q: str, current_user: User = Security(get_current_active_user, scopes=["albums.list"])):
 
     output = {"results": []}
@@ -50,7 +51,7 @@ async def album_find(q: str, current_user: User = Security(get_current_active_us
 
     return UJSONResponse(output)
 
-@app.patch("/albums/{id}", response_class=UJSONResponse, description="Modify album's name or title by id")
+@app.patch("/albums/{id}", response_class=UJSONResponse, response_model=AlbumModified, description="Modify album's name or title by id")
 async def album_patch(id: str, name: Union[str, None] = None, title: Union[str, None] = None, current_user: User = Security(get_current_active_user, scopes=["albums.write"])):
 
     try:
@@ -88,7 +89,7 @@ async def album_patch(id: str, name: Union[str, None] = None, title: Union[str,
         }
     )
 
-@app.put("/albums/{id}", response_class=UJSONResponse, description="Modify album's name and title by id")
+@app.put("/albums/{id}", response_class=UJSONResponse, response_model=AlbumModified, description="Modify album's name and title by id")
 async def album_put(id: str, name: str, title: str, current_user: User = Security(get_current_active_user, scopes=["albums.write"])):
 
     try:

extensions/photos.py

@@ -4,6 +4,7 @@ from secrets import token_urlsafe
 from magic import Magic
 from datetime import datetime
 from os import makedirs, path, remove
+from classes.models import Photo, SearchResults
 from modules.hasher import get_phash, get_duplicates
 from modules.security import User, get_current_active_user
 from modules.app import app
@@ -15,7 +16,7 @@ from fastapi import HTTPException, UploadFile, Security
 from fastapi.responses import UJSONResponse, Response
 from starlette.status import HTTP_204_NO_CONTENT, HTTP_400_BAD_REQUEST, HTTP_401_UNAUTHORIZED, HTTP_404_NOT_FOUND, HTTP_409_CONFLICT
 
-@app.post("/albums/{album}/photos", response_class=UJSONResponse, description="Upload a photo to album")
+@app.post("/albums/{album}/photos", response_class=UJSONResponse, response_model=Photo, description="Upload a photo to album")
 async def photo_upload(file: UploadFile, album: str, ignore_duplicates: bool = False, current_user: User = Security(get_current_active_user, scopes=["photos.write"])):
 
     if col_albums.find_one( {"user": current_user.user, "name": album} ) is None:
@@ -91,7 +92,7 @@ async def photo_delete(id: str, current_user: User = Security(get_current_active
 
     return Response(status_code=HTTP_204_NO_CONTENT)
 
-@app.get("/albums/{album}/photos", response_class=UJSONResponse, description="Find a photo by filename")
+@app.get("/albums/{album}/photos", response_class=UJSONResponse, response_model=SearchResults, description="Find a photo by filename")
 async def photo_find(q: str, album: str, page: int = 1, page_size: int = 100, current_user: User = Security(get_current_active_user, scopes=["photos.list"])):
 
     if col_albums.find_one( {"user": current_user.user, "name": album} ) is None:
@@ -111,13 +112,12 @@ async def photo_find(q: str, album: str, page: int = 1, page_size: int = 100, cu
         token = str(token_urlsafe(32))
         col_tokens.insert_one( {"token": token, "query": q, "album": album, "page": page+1, "page_size": page_size, "user": pickle.dumps(current_user)} )
         output["next_page"] = f"/albums/{album}/photos/token?token={token}" # type: ignore
-
-    with open("something.txt", "w", encoding="utf-8") as f:
-        f.write(pickle.loads(pickle.dumps(current_user)).user)
+    else:
+        output["next_page"] = None # type: ignore
 
     return UJSONResponse(output)
 
-@app.get("/albums/{album}/photos/token", response_class=UJSONResponse, description="Find a photo by token")
+@app.get("/albums/{album}/photos/token", response_class=UJSONResponse, response_model=SearchResults, description="Find a photo by token")
 async def photo_find_token(token: str):
 
     found_record = col_tokens.find_one( {"token": token} )