diff --git a/migrations/20250219073939_cleanup_episodes.py b/migrations/20250219073939_cleanup_episodes.py new file mode 100644 index 0000000..5a89730 --- /dev/null +++ b/migrations/20250219073939_cleanup_episodes.py @@ -0,0 +1,201 @@ +from datetime import datetime, timezone +from typing import List, Optional, Tuple +from dataclasses import dataclass + +from beanie import Document, free_fall_migration +from pydantic import BaseModel, Field +import PTT +from tqdm.asyncio import tqdm +import logging + +# Configure logging +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - [%(name)s] - %(message)s" +) +logger = logging.getLogger(__name__) + + +class EpisodeFile(BaseModel): + season_number: int + episode_number: int + filename: Optional[str] = None + + +class TorrentStreams(Document): + id: str + meta_id: str + torrent_name: str + size: int + episode_files: List[EpisodeFile] = Field(default_factory=list) + updated_at: datetime | None = Field( + default_factory=lambda: datetime.now(tz=timezone.utc) + ) + + class Settings: + name = "TorrentStreams" + + +@dataclass +class EpisodeDetail: + season_number: int + episode_number: int + + def to_tuple(self) -> Tuple[int, int]: + return (self.season_number, self.episode_number) + + def to_episode_file(self) -> EpisodeFile: + return EpisodeFile( + season_number=self.season_number, episode_number=self.episode_number + ) + + +class EpisodeProcessor: + def __init__(self, stream: TorrentStreams): + self.stream = stream + self.parsed_title = PTT.parse_title(stream.torrent_name) + + def _parse_seasons_episodes(self) -> Tuple[List[int], List[int]]: + seasons = self.parsed_title.get("seasons", []) + episodes = self.parsed_title.get("episodes", []) + + if not seasons and not episodes: + seasons = [1] + + return seasons, episodes + + def generate_episode_details(self) -> List[EpisodeDetail]: + seasons, episodes = self._parse_seasons_episodes() + + if len(seasons) == 1 and episodes: + return [EpisodeDetail(seasons[0], episode) for episode in episodes] + elif len(seasons) == 1: + return [EpisodeDetail(seasons[0], 1)] + elif len(seasons) > 1: + return [EpisodeDetail(season, 1) for season in seasons] + elif episodes: + return [EpisodeDetail(1, episode) for episode in episodes] + else: + return [EpisodeDetail(1, 1)] + + def process_episodes(self) -> List[EpisodeFile]: + existing_episodes = self.stream.episode_files + episode_details = self.generate_episode_details() + + # Early return if no episode details generated + if not episode_details: + return existing_episodes + + exist_any_filename = any(ep.filename for ep in existing_episodes) + + if exist_any_filename: + return self._process_with_existing_filenames( + existing_episodes, episode_details + ) + else: + return self._process_without_filenames(existing_episodes, episode_details) + + def _process_with_existing_filenames( + self, existing_episodes: List[EpisodeFile], episode_details: List[EpisodeDetail] + ) -> List[EpisodeFile]: + cleaned_episodes = [ep for ep in existing_episodes if ep.filename] + cleaned_episodes_keys = { + (ep.season_number, ep.episode_number) for ep in cleaned_episodes + } + + # Log removed episodes + removed_episodes = { + (ep.season_number, ep.episode_number) + for ep in existing_episodes + if not ep.filename + } + if removed_episodes: + logger.debug( + f"Removed empty filename episodes for stream {self.stream.id} " + f"with {removed_episodes}" + ) + + # Add missing episodes + new_episodes = [ + detail.to_episode_file() + for detail in episode_details + if detail.to_tuple() not in cleaned_episodes_keys + ] + cleaned_episodes.extend(new_episodes) + return cleaned_episodes + + def _process_without_filenames( + self, existing_episodes: List[EpisodeFile], episode_details: List[EpisodeDetail] + ) -> List[EpisodeFile]: + existing_episodes_keys = { + (ep.season_number, ep.episode_number) for ep in existing_episodes + } + + new_episodes = [detail.to_episode_file() for detail in episode_details] + + # Log changes + changes = [ + detail + for detail in episode_details + if detail.to_tuple() not in existing_episodes_keys + ] + if changes: + logger.debug(f"Torrent {self.stream.id} Change episode details: {changes}") + + return new_episodes + + +class Forward: + @free_fall_migration(document_models=[TorrentStreams]) + async def cleanup_episodes(self, session): + query = { + "episode_files": { + "$elemMatch": { + "season_number": {"$exists": True}, + "$or": [{"filename": None}, {"filename": {"$exists": False}}], + } + }, + } + + streams = TorrentStreams.find(query) + count = await streams.count() + logger.info(f"Found {count} streams with empty filename episodes") + + with tqdm(total=count, desc="Processing streams") as pbar: + async for stream in streams: + try: + processor = EpisodeProcessor(stream) + cleaned_episodes = processor.process_episodes() + + # Skip update if no changes + if cleaned_episodes == stream.episode_files: + pbar.update(1) + continue + + stream.episode_files = cleaned_episodes + new_keys = [ + (ep.season_number, ep.episode_number) + for ep in stream.episode_files + ] + + logger.debug( + f"Updated episodes for stream {stream.id} '{stream.torrent_name}' " + f"with {new_keys}" + ) + + stream.updated_at = datetime.now(tz=timezone.utc) + await stream.save() + + except Exception as e: + logger.error( + f"Error processing stream {stream.id}: {str(e)}", exc_info=True + ) + continue + + pbar.update(1) + + logger.info("Migration completed successfully") + + +class Backward: + + pass