diff --git a/src/program/media/__init__.py b/src/program/media/__init__.py index 271ab569..1bf3c795 100644 --- a/src/program/media/__init__.py +++ b/src/program/media/__init__.py @@ -1,2 +1,2 @@ -from .item import Episode, MediaItem, Movie, Season, Show # noqa +from .item import Episode, MediaItem, Movie, Season, Show, ShowMediaType, MovieMediaType, MediaType # noqa from .state import States # noqa diff --git a/src/program/media/item.py b/src/program/media/item.py index 4d9b9e7b..ae34fee3 100644 --- a/src/program/media/item.py +++ b/src/program/media/item.py @@ -1,5 +1,6 @@ """MediaItem class""" from datetime import datetime +from enum import Enum from pathlib import Path from typing import List, Optional, Self @@ -17,6 +18,22 @@ from ..db.db_functions import blacklist_stream, reset_streams from .stream import Stream +class ShowMediaType(Enum): + """Show media types""" + Show = "show" + Season = "season" + Episode = "episode" + +class MovieMediaType(Enum): + """Media types""" + Movie = "movie" + +class MediaType(Enum): + """Combined media types""" + Show = ShowMediaType.Show.value + Season = ShowMediaType.Season.value + Episode = ShowMediaType.Episode.value + Movie = MovieMediaType.Movie.value class MediaItem(db.Model): """MediaItem class""" @@ -406,7 +423,7 @@ def copy(self, other): return self def __init__(self, item): - self.type = "movie" + self.type = MovieMediaType.Movie.value self.file = item.get("file", None) super().__init__(item) @@ -429,7 +446,7 @@ class Show(MediaItem): def __init__(self, item): super().__init__(item) - self.type = "show" + self.type = ShowMediaType.Show.value self.locations = item.get("locations", []) self.seasons: list[Season] = item.get("seasons", []) self.propagate_attributes_to_childs() @@ -543,7 +560,7 @@ def store_state(self, given_state: States = None) -> None: super().store_state(given_state) def __init__(self, item): - self.type = "season" + self.type = ShowMediaType.Season.value self.number = item.get("number", None) self.episodes: list[Episode] = item.get("episodes", []) self.parent = item.get("parent", None) @@ -643,7 +660,7 @@ class Episode(MediaItem): } def __init__(self, item): - self.type = "episode" + self.type = ShowMediaType.Episode.value self.number = item.get("number", None) self.file = item.get("file", None) super().__init__(item) diff --git a/src/program/services/downloaders/__init__.py b/src/program/services/downloaders/__init__.py index c80a1dc7..a01988f4 100644 --- a/src/program/services/downloaders/__init__.py +++ b/src/program/services/downloaders/__init__.py @@ -1,15 +1,15 @@ -from concurrent.futures import CancelledError, ThreadPoolExecutor, as_completed - from loguru import logger from program.media.item import MediaItem from program.media.state import States from program.media.stream import Stream from program.settings.manager import settings_manager -from program.services.downloaders.shared import filesize_is_acceptable_movie, filesize_is_acceptable_show +from program.services.downloaders.shared import filesize_is_acceptable from .alldebrid import AllDebridDownloader from .realdebrid import RealDebridDownloader +from ...media import MovieMediaType, ShowMediaType + # from .torbox import TorBoxDownloader @@ -91,24 +91,23 @@ def update_item_attributes(self, item: MediaItem, info, container) -> bool: item = item container = container for file in container.values(): - if item.type == "movie" and self.service.file_finder.container_file_matches_movie(file) and filesize_is_acceptable_movie(file[self.service.file_finder.filesize_attr]): + if not self.validate_filesize(item, info, file): + continue + if item.type == MovieMediaType.Movie.value and self.service.file_finder.container_file_matches_movie(file): item.file = file[self.service.file_finder.filename_attr] item.folder = info["filename"] item.alternative_folder = info["original_filename"] item.active_stream = {"infohash": info["hash"], "id": info["id"]} found = True break - else: - logger.debug(f"Deleting {info['id']} because it doesn't match the item type or size") - self.delete_torrent(info["id"]) - if item.type in ["show", "season", "episode"]: + if item.type in (media_type.value for media_type in ShowMediaType): show = item - if item.type == "season": + if item.type == ShowMediaType.Season.value: show = item.parent - elif item.type == "episode": + elif item.type == ShowMediaType.Episode.value: show = item.parent.parent file_season, file_episodes = self.service.file_finder.container_file_matches_episode(file) - if file_season and file_episodes and filesize_is_acceptable_show(file[self.service.file_finder.filesize_attr]): + if file_season and file_episodes: season = next((season for season in show.seasons if season.number == file_season), None) for file_episode in file_episodes: episode = next((episode for episode in season.episodes if episode.number == file_episode), None) @@ -118,9 +117,24 @@ def update_item_attributes(self, item: MediaItem, info, container) -> bool: episode.alternative_folder = info["original_filename"] episode.active_stream = {"infohash": info["hash"], "id": info["id"]} # We have to make sure the episode is correct if item is an episode - if item.type != "episode" or (item.type == "episode" and episode.number == item.number): + if item.type != ShowMediaType.Episode.value or (item.type == ShowMediaType.Episode.value and episode.number == item.number): found = True - else: - logger.debug(f"Deleting {info['id']} because it doesn't match the item type or size") - self.delete_torrent(info["id"]) - return found \ No newline at end of file + return found + + def validate_filesize(self, item, info, file) -> bool: + item_media_type = self._get_item_media_type(item) + if filesize_is_acceptable(file[self.service.file_finder.filesize_attr], item_media_type): + return True + torrent_id = info.get("id", None) + if not torrent_id: + logger.warning(f"Couldn't find torrent id to delete") + return False + logger.debug(f"Deleting {torrent_id} because it doesn't match the item type or size") + self.service.delete_torrent(torrent_id) + return False + + @staticmethod + def _get_item_media_type(item): + if item.type in (media_type.value for media_type in ShowMediaType): + return ShowMediaType.Show.value + return MovieMediaType.Movie.value \ No newline at end of file diff --git a/src/program/services/downloaders/shared.py b/src/program/services/downloaders/shared.py index a187a80f..7fa9efe3 100644 --- a/src/program/services/downloaders/shared.py +++ b/src/program/services/downloaders/shared.py @@ -5,6 +5,7 @@ from loguru import logger from RTN import parse +from program.media import MovieMediaType, ShowMediaType from program.settings.manager import settings_manager DEFAULT_VIDEO_EXTENSIONS = ["mp4", "mkv", "avi"] @@ -123,7 +124,6 @@ def hash_from_uri(magnet_uri: str) -> str: max_movie_filesize = settings_manager.settings.downloaders.movie_filesize_mb_max min_episode_filesize = settings_manager.settings.downloaders.episode_filesize_mb_min max_episode_filesize = settings_manager.settings.downloaders.episode_filesize_mb_max -are_filesizes_valid = False def _validate_filesize_setting(value: int, setting_name: str) -> bool: """Validate a single filesize setting.""" @@ -156,7 +156,7 @@ def _convert_to_bytes(size_mb: int) -> int: def _get_size_limits(media_type: str) -> Tuple[int, int]: """Get min and max size limits in MB for given media type.""" settings = settings_manager.settings.downloaders - if media_type == "movie": + if media_type == MovieMediaType.Movie.value: return (settings.movie_filesize_mb_min, settings.movie_filesize_mb_max) return (settings.episode_filesize_mb_min, settings.episode_filesize_mb_max) @@ -185,8 +185,5 @@ def _validate_filesize(filesize: int, media_type: str) -> bool: return is_acceptable -def filesize_is_acceptable_movie(filesize: int) -> bool: - return _validate_filesize(filesize, "movie") - -def filesize_is_acceptable_show(filesize: int) -> bool: - return _validate_filesize(filesize, "show") \ No newline at end of file +def filesize_is_acceptable(filesize: int, media_type: str) -> bool: + return _validate_filesize(filesize, media_type) \ No newline at end of file