from collections import namedtuple from difflib import SequenceMatcher import io import logging import os import re import zipfile from guessit import guessit import rarfile from subliminal.subtitle import fix_line_ending from subliminal_patch.core import Episode from subliminal_patch.subtitle import guess_matches from ._agent_list import FIRST_THOUSAND_OR_SO_USER_AGENTS logger = logging.getLogger(__name__) _MatchingSub = namedtuple("_MatchingSub", ("file", "priority")) def _get_matching_sub(sub_names, forced=False, episode=None, episode_title=None): guess_options = {"single_value": True} if episode is not None: guess_options["type"] = "episode" # type: ignore matching_subs = [] for sub_name in sub_names: if not forced and os.path.splitext(sub_name.lower())[0].endswith("forced"): logger.debug("Ignoring forced subtitle: %s", sub_name) continue # If it's a movie then get the first subtitle if episode is None and episode_title is None: logger.debug("Movie subtitle found: %s", sub_name) matching_subs.append(_MatchingSub(sub_name, 2)) break guess = guessit(sub_name, options=guess_options) matched_episode_num = guess.get("episode") if matched_episode_num: logger.debug("No episode number found in file: %s", sub_name) matched_title = None if episode_title is not None: matched_title = _analize_sub_name(sub_name, episode_title) if episode == matched_episode_num: logger.debug("Episode matched from number: %s", sub_name) matching_subs.append(_MatchingSub(sub_name, 2)) elif matched_title: matching_subs.append(_MatchingSub(sub_name, 1)) else: logger.debug("Ignoring incorrect episode: '%s'", sub_name) if matching_subs: matching_subs.sort(key=lambda x: x.priority, reverse=True) logger.debug("Matches: %s", matching_subs) return matching_subs[0].file else: logger.debug("Nothing matched") return None def _analize_sub_name(sub_name: str, title_): titles = re.split(r"[.-]", os.path.splitext(sub_name)[0]) for title in titles: ratio = SequenceMatcher(None, title, title_).ratio() if ratio > 0.85: logger.debug( "Episode title matched: '%s' -> '%s' [%s]", title, sub_name, ratio ) return True logger.debug("No episode title matched from file") return False def get_subtitle_from_archive( archive, forced=False, episode=None, get_first_subtitle=False, **kwargs ): "Get subtitle from Rarfile/Zipfile object. Return None if nothing is found." subs_in_archive = [ name for name in archive.namelist() if name.endswith((".srt", ".sub", ".ssa", ".ass")) ] if not subs_in_archive: logger.info("No subtitles found in archive") return None logger.debug("Subtitles in archive: %s", subs_in_archive) if len(subs_in_archive) == 1 or get_first_subtitle: logger.debug("Getting first subtitle in archive: %s", subs_in_archive) return fix_line_ending(archive.read(subs_in_archive[0])) matching_sub = _get_matching_sub(subs_in_archive, forced, episode, **kwargs) if matching_sub is not None: logger.info("Using %s from archive", matching_sub) return fix_line_ending(archive.read(matching_sub)) logger.debug("No subtitle found in archive") return None def is_episode(content): return "episode" in guessit(content, {"type": "episode"}) def get_archive_from_bytes(content: bytes): """Get RarFile/ZipFile object from bytes. Return None is something else is found.""" # open the archive archive_stream = io.BytesIO(content) if rarfile.is_rarfile(archive_stream): logger.debug("Identified rar archive") return rarfile.RarFile(archive_stream) elif zipfile.is_zipfile(archive_stream): logger.debug("Identified zip archive") return zipfile.ZipFile(archive_stream) logger.debug("Unknown compression format") return None def update_matches(matches, video, release_info: str, **guessit_options): "Update matches set from release info string. New lines are iterated." guessit_options["type"] = "episode" if isinstance(video, Episode) else "movie" logger.debug("Guessit options to update matches: %s", guessit_options) for release in release_info.split("\n"): logger.debug("Updating matches from release info: %s", release) matches |= guess_matches(video, guessit(release.strip(), guessit_options)) logger.debug("New matches: %s", matches) return matches