From 44f047e7e00c58628fa0669f1630b80f8bbe936e Mon Sep 17 00:00:00 2001
From: Spoked <dreu.lavelle@gmail.com>
Date: Sun, 1 Dec 2024 01:19:47 -0500
Subject: [PATCH] fix: add strong typed response to scrape api endpoint

---
 src/program/services/downloaders/models.py | 19 -------------------
 src/routers/secure/scrape.py               | 14 ++++++++------
 2 files changed, 8 insertions(+), 25 deletions(-)

diff --git a/src/program/services/downloaders/models.py b/src/program/services/downloaders/models.py
index 485eb7ea..5d09181b 100644
--- a/src/program/services/downloaders/models.py
+++ b/src/program/services/downloaders/models.py
@@ -71,13 +71,6 @@ def create(
 
         return cls(filename=filename, filesize=filesize_bytes, file_id=file_id)
 
-    def to_dict(self) -> dict:
-        """Convert the DebridFile to a dictionary"""
-        return {
-            "file_id": self.file_id,
-            "filename": self.filename,
-            "filesize": self.filesize
-        }
 
 class ParsedFileData(BaseModel):
     """Represents a parsed file from a filename"""
@@ -101,15 +94,6 @@ def file_ids(self) -> List[int]:
         """Get the file ids of the cached files"""
         return [file.file_id for file in self.files if file.file_id is not None]
 
-    def to_dict(self) -> dict:
-        """Convert the TorrentContainer to a dictionary"""
-        return {
-            file.file_id or str(i): {
-                "filename": file.filename,
-                "filesize": file.filesize
-            }
-            for i, file in enumerate(self.files)
-        }
 
 class TorrentInfo(BaseModel):
     """Torrent information from a debrid service"""
@@ -130,9 +114,6 @@ def size_mb(self) -> float:
         """Convert bytes to megabytes"""
         return self.bytes / 1_000_000
 
-    def to_dict(self) -> dict:
-        """Convert the TorrentInfo to a dictionary"""
-        return self.model_dump()
 
 class DownloadedTorrent(BaseModel):
     """Represents the result of a download operation"""
diff --git a/src/routers/secure/scrape.py b/src/routers/secure/scrape.py
index 8e794e40..3df265c6 100644
--- a/src/routers/secure/scrape.py
+++ b/src/routers/secure/scrape.py
@@ -39,8 +39,8 @@ class StartSessionResponse(BaseModel):
     message: str
     session_id: str
     torrent_id: str
-    torrent_info: dict
-    containers: Optional[List[dict]]
+    torrent_info: TorrentInfo
+    containers: Optional[List[TorrentContainer]]
     expires_at: str
 
 class SelectFilesResponse(BaseModel):
@@ -276,7 +276,7 @@ def get_info_hash(magnet: str) -> str:
     session = session_manager.create_session(item_id or imdb_id, info_hash)
 
     try:
-        torrent_id: Union[int, str] = downloader.add_torrent(info_hash)
+        torrent_id: str = downloader.add_torrent(info_hash)
         torrent_info: TorrentInfo = downloader.get_torrent_info(torrent_id)
         container: Optional[TorrentContainer] = downloader.get_instant_availability(info_hash, item.type)
         session_manager.update_session(session.id, torrent_id=torrent_id, torrent_info=torrent_info, containers=container)
@@ -284,15 +284,17 @@ def get_info_hash(magnet: str) -> str:
         background_tasks.add_task(session_manager.abort_session, session.id)
         raise HTTPException(status_code=500, detail=str(e))
 
-    return {
+    data = {
         "message": "Started manual scraping session",
         "session_id": session.id,
         "torrent_id": torrent_id,
-        "torrent_info": torrent_info.to_dict(),
-        "containers": [container.to_dict()] if container else None,
+        "torrent_info": torrent_info,
+        "containers": [container] if container else None,
         "expires_at": session.expires_at.isoformat()
     }
 
+    return StartSessionResponse(**data)
+
 @router.post(
     "/scrape/select_files/{session_id}",
     summary="Select files for torrent id, for this to be instant it requires files to be one of /manual/instant_availability response containers",