diff --git a/src/program/services/downloaders/models.py b/src/program/services/downloaders/models.py index c90dbe97..ddcd0864 100644 --- a/src/program/services/downloaders/models.py +++ b/src/program/services/downloaders/models.py @@ -47,21 +47,37 @@ class DebridFile(BaseModel): filesize: Optional[int] = Field(default=None) @classmethod - def create(cls, filename: str, filesize_bytes: int, filetype: Literal["movie", "episode"], file_id: Optional[int] = None) -> Optional["DebridFile"]: + def create( + cls, + filename: str, + filesize_bytes: int, + filetype: Literal["movie", "show", "season", "episode"], + file_id: Optional[int] = None, + limit_filesize: bool = True + + ) -> Optional["DebridFile"]: """Factory method to validate and create a DebridFile""" if not any(filename.endswith(ext) for ext in VIDEO_EXTENSIONS) and not "sample" in filename.lower(): return None - - filesize_mb = filesize_bytes / 1_000_000 - if filetype == "movie": - if not (FILESIZE_MOVIE_CONSTRAINT[0] <= filesize_mb <= FILESIZE_MOVIE_CONSTRAINT[1]): - return None - elif filetype == "episode": - if not (FILESIZE_EPISODE_CONSTRAINT[0] <= filesize_mb <= FILESIZE_EPISODE_CONSTRAINT[1]): - return None + + if limit_filesize: + filesize_mb = filesize_bytes / 1_000_000 + if filetype == "movie": + if not (FILESIZE_MOVIE_CONSTRAINT[0] <= filesize_mb <= FILESIZE_MOVIE_CONSTRAINT[1]): + return None + elif filetype in ["show", "season", "episode"]: + if not (FILESIZE_EPISODE_CONSTRAINT[0] <= filesize_mb <= FILESIZE_EPISODE_CONSTRAINT[1]): + return None 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""" @@ -85,6 +101,12 @@ 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 { + "infohash": self.infohash, + "files": [file.to_dict() for file in self.files] + } class TorrentInfo(BaseModel): """Torrent information from a debrid service""" @@ -105,6 +127,13 @@ 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""" + files = [file.to_dict() for file in self.files] + return { + **self.model_dump(), + "files": files + } class DownloadedTorrent(BaseModel): """Represents the result of a download operation""" diff --git a/src/program/services/downloaders/shared.py b/src/program/services/downloaders/shared.py index c71ee9f4..1fd96c4c 100644 --- a/src/program/services/downloaders/shared.py +++ b/src/program/services/downloaders/shared.py @@ -1,6 +1,6 @@ from abc import ABC, abstractmethod from datetime import datetime -from typing import Optional +from typing import Optional, Union from RTN import ParsedData, parse @@ -41,7 +41,7 @@ def get_instant_availability(self, infohash: str, item_type: str) -> Optional[To pass @abstractmethod - def add_torrent(self, infohash: str) -> int: + def add_torrent(self, infohash: str) -> Union[int, str]: """ Add a torrent and return its information @@ -49,17 +49,21 @@ def add_torrent(self, infohash: str) -> int: infohash: The hash of the torrent to add Returns: - str: The ID of the added torrent + Union[int, str]: The ID of the added torrent + + Notes: + The return type changes depending on the downloader """ pass @abstractmethod - def select_files(self, request: list[int]) -> None: + def select_files(self, torrent_id: str, file_ids: list[int]) -> None: """ Select which files to download from the torrent Args: - request: File selection details including torrent ID and file IDs + torrent_id: ID of the torrent to select files for + file_ids: IDs of the files to select """ pass diff --git a/src/routers/secure/scrape.py b/src/routers/secure/scrape.py index 132ef2fc..e801071b 100644 --- a/src/routers/secure/scrape.py +++ b/src/routers/secure/scrape.py @@ -19,6 +19,7 @@ from program.services.scrapers import Scraping from program.services.scrapers.shared import rtn from program.types import Event +from program.services.downloaders.models import TorrentContainer, TorrentInfo class Stream(BaseModel): @@ -218,13 +219,10 @@ def scrape_item(request: Request, id: str) -> ScrapeItemResponse: .unique() .scalar_one_or_none() ) - streams = scraper.scrape(item) - stream_containers = downloader.get_instant_availability([stream for stream in streams.keys()]) - for stream in streams.keys(): - if len(stream_containers.get(stream, [])) > 0: - streams[stream].is_cached = True - else: - streams[stream].is_cached = False + streams: Dict[str, Stream] = scraper.scrape(item) + for stream in streams.values(): + container = downloader.get_instant_availability(stream.infohash, item.type) + stream.is_cached = bool(container and container.cached) log_string = item.log_string return { @@ -278,10 +276,10 @@ def get_info_hash(magnet: str) -> str: session = session_manager.create_session(item_id or imdb_id, info_hash) try: - torrent_id = downloader.add_torrent(info_hash) - torrent_info = downloader.get_torrent_info(torrent_id) - containers = downloader.get_instant_availability([session.magnet]).get(session.magnet, None) - session_manager.update_session(session.id, torrent_id=torrent_id, torrent_info=torrent_info, containers=containers) + torrent_id: Union[int, 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)) @@ -290,8 +288,8 @@ def get_info_hash(magnet: str) -> str: "message": "Started manual scraping session", "session_id": session.id, "torrent_id": torrent_id, - "torrent_info": torrent_info, - "containers": containers, + "torrent_info": torrent_info.model_dump_json() if torrent_info else None, + "containers": [container.model_dump_json()] if container else None, "expires_at": session.expires_at.isoformat() } @@ -307,7 +305,7 @@ def manual_select_files(request: Request, session_id, files: Container) -> Selec raise HTTPException(status_code=404, detail="Session not found or expired") if not session.torrent_id: session_manager.abort_session(session_id) - raise HTTPException(status_code=500, detail="") + raise HTTPException(status_code=500, detail="No torrent ID found") download_type = "uncached" if files.model_dump() in session.containers: