view index fix

src/endpoints/view.py

@@ -1,6 +1,8 @@
 from fastapi import APIRouter
 
 from src.libs.s3 import s3resource
+from src.models.view import ViewResult, ViewImage
+from src.models.main import User, Method
 from uuid import UUID
 
 router = APIRouter(
@@ -16,10 +18,32 @@ async def status():
 
 
 @router.get("/")
-async def view(user: UUID, index: int | None = None):
+async def view(user: UUID, index: int | None = None) -> ViewResult:
     s3 = s3resource()
     bucket = s3.Bucket('pimthaigans-image-container')
+
     objs = list(bucket.objects.filter(Prefix=f'{user}/'))
+    path_objs = [obj.key for obj in objs]
+
+    if len(path_objs) == 0:
+        return ViewResult(user=User(uuid=user), method=Method.index, result=None)
+
     if index:
-        return objs[index].key
-    return [obj.key for obj in objs]
+        result: ViewResult = await view_index(User(uuid=user), index, path_objs)
+        return result
+
+    result: ViewResult = await view_all(User(uuid=user), path_objs)
+    return result
+
+async def view_index(user: User, index: int, path_objs) -> ViewResult:
+    imgs = [obj for obj in path_objs if obj.endswith(f'{str(index).zfill(2)}.png')]
+    if len(imgs) > 0:
+        return ViewResult(user=user, method=Method.index, result=[ViewImage(index=index, image_url=imgs[0])])
+
+    return ViewResult(user=user, method=Method.index, result=None)
+
+async def view_all(user: User, path_objs) -> ViewResult:
+    for index in range(0, 88):
+        imgs = [obj for obj in path_objs if obj.endswith(f'{str(index).zfill(2)}.png')]
+        if len(imgs) > 0:
+            return ViewResult(user=user, method=Method.all, result=[ViewImage(index=index, image_url=imgs[0])])

src/models/view.py

@@ -0,0 +1,15 @@
+from typing import List
+
+from pydantic import BaseModel
+
+from src.models.main import User, Method
+
+
+class ViewImage(BaseModel):
+    index: int
+    image_url: str
+
+class ViewResult(BaseModel):
+    user: User
+    method: Method
+    result: List[ViewImage] | None