From 653688b3763dab10740c005c4fc98988579912a6 Mon Sep 17 00:00:00 2001 From: Profitroll <47523801+profitrollgame@users.noreply.github.com> Date: Tue, 20 Dec 2022 17:07:48 +0100 Subject: [PATCH] Added models where possible --- classes/models.py | 21 +++++++++++++++++++++ extensions/albums.py | 9 +++++---- extensions/photos.py | 12 ++++++------ 3 files changed, 32 insertions(+), 10 deletions(-) create mode 100644 classes/models.py diff --git a/classes/models.py b/classes/models.py new file mode 100644 index 0000000..15a2883 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/extensions/albums.py b/extensions/albums.py index 602bf0f..5d3f0ce 100644 --- a/extensions/albums.py +++ b/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: diff --git a/extensions/photos.py b/extensions/photos.py index e47ec75..10cf68d 100644 --- a/extensions/photos.py +++ b/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} )