Skip to content

Commit

Permalink
fix: add strong typed response to scrape api endpoint
Browse files Browse the repository at this point in the history
  • Loading branch information
dreulavelle committed Dec 5, 2024
1 parent 346b352 commit 44f047e
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 25 deletions.
19 changes: 0 additions & 19 deletions src/program/services/downloaders/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"""
Expand All @@ -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"""
Expand All @@ -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"""
Expand Down
14 changes: 8 additions & 6 deletions src/routers/secure/scrape.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -276,23 +276,25 @@ 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)
except Exception as e:
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",
Expand Down

0 comments on commit 44f047e

Please sign in to comment.