From 848836c5bb66ae4baa5edac3decfe65de4cb9e62 Mon Sep 17 00:00:00 2001 From: vitiko98 Date: Mon, 23 May 2022 22:43:34 -0400 Subject: [PATCH] Subf2m provider: add season pack support --- libs/subliminal_patch/providers/subf2m.py | 81 +++++++++++++++++------ tests/subliminal_patch/test_subf2m.py | 3 +- 2 files changed, 60 insertions(+), 24 deletions(-) diff --git a/libs/subliminal_patch/providers/subf2m.py b/libs/subliminal_patch/providers/subf2m.py index 689a4aac7..8ce43ca5f 100644 --- a/libs/subliminal_patch/providers/subf2m.py +++ b/libs/subliminal_patch/providers/subf2m.py @@ -2,18 +2,17 @@ import logging -from requests import Session from bs4 import BeautifulSoup as bso - -from subliminal_patch.exceptions import APIThrottled +from requests import Session from subliminal_patch.core import Episode from subliminal_patch.core import Movie +from subliminal_patch.exceptions import APIThrottled from subliminal_patch.providers import Provider -from subliminal_patch.subtitle import Subtitle from subliminal_patch.providers.utils import get_archive_from_bytes from subliminal_patch.providers.utils import get_subtitle_from_archive +from subliminal_patch.providers.utils import is_episode from subliminal_patch.providers.utils import update_matches - +from subliminal_patch.subtitle import Subtitle from subzero.language import Language logger = logging.getLogger(__name__) @@ -23,11 +22,17 @@ class Subf2mSubtitle(Subtitle): provider_name = "subf2m" hash_verifiable = False - def __init__(self, language, page_link, release_info, pre_matches): + def __init__(self, language, page_link, release_info, episode_number=None): super().__init__(language, page_link=page_link) self.release_info = release_info - self._matches = set(pre_matches) + self.episode_number = episode_number + + self._matches = set( + ("title", "year") + if episode_number is None + else ("title", "series", "season", "episode") + ) def get_matches(self, video): update_matches(self._matches, video, self.release_info) @@ -154,7 +159,7 @@ class Subf2mProvider(Provider): subtitles = [] for item in soup.select("li.item"): - subtitle = _get_subtitle_from_item(item, language, Movie) + subtitle = _get_subtitle_from_item(item, language) if subtitle is None: continue @@ -165,19 +170,39 @@ class Subf2mProvider(Provider): def _find_episode_subtitles(self, path, season, episode, language): # TODO: add season packs support? - soup = self._get_subtitle_page_soup(path, language) - expected_substring = f"s{season:02}e{episode:02}".lower() + + season_pack_substrings = _get_season_pack_substrings(season) + subtitles = [] for item in soup.select("li.item"): - if expected_substring in item.text.lower(): - subtitle = _get_subtitle_from_item(item, language, Episode) - if subtitle is None: + valid_item = None + + if not item.text.strip(): + continue + + if f"s{season:02}e{episode:02}" in item.text.lower(): + valid_item = item + + elif any(sp in item.text.lower() for sp in season_pack_substrings): + logger.debug("Possible season pack found") + if is_episode(item.text): + logger.debug("It's an episode: %s", " ".join(item.text.split())) continue - logger.debug("Found subtitle: %s", subtitle) - subtitles.append(subtitle) + logger.debug("Season pack found: %s", " ".join(item.text.split())) + valid_item = item + + if valid_item is None: + continue + + subtitle = _get_subtitle_from_item(item, language, episode) + if subtitle is None: + continue + + logger.debug("Found subtitle: %s", subtitle) + subtitles.append(subtitle) return subtitles @@ -235,16 +260,28 @@ class Subf2mProvider(Provider): if archive is None: raise APIThrottled(f"Invalid archive: {subtitle.page_link}") - subtitle.content = get_subtitle_from_archive(archive, get_first_subtitle=True) + subtitle.content = get_subtitle_from_archive( + archive, episode=subtitle.episode_number + ) -_types_map = { - Movie: ("title", "year"), - Episode: ("title", "series", "season", "episode"), -} +def _get_season_pack_substrings(season): + season_pack_substrings = [ + f"season {season:02}", + f"season {season}", + f" s{season:02}", + f" s{season}", + ] + + try: + season_pack_substrings.append(f"{_SEASONS[season - 1]} season") + except IndexError: + pass + + return season_pack_substrings -def _get_subtitle_from_item(item, language, type): +def _get_subtitle_from_item(item, language, episode_number=None): release_info = [ rel.text.strip() for rel in item.find("ul", {"class": "scrolllist"}) ] @@ -263,4 +300,4 @@ def _get_subtitle_from_item(item, language, type): logger.debug("Couldn't get path: %s", item) return None - return Subf2mSubtitle(language, _BASE_URL + path, release_info, _types_map[type]) + return Subf2mSubtitle(language, _BASE_URL + path, release_info, episode_number) diff --git a/tests/subliminal_patch/test_subf2m.py b/tests/subliminal_patch/test_subf2m.py index 4bac1c030..49651f3d5 100644 --- a/tests/subliminal_patch/test_subf2m.py +++ b/tests/subliminal_patch/test_subf2m.py @@ -54,7 +54,6 @@ def subtitle(): Language.fromalpha3b("per"), "https://subf2m.co/subtitles/dune-2021/farsi_persian/2604701", release_info, - {"title", "year"}, ) @@ -64,7 +63,7 @@ def subtitle_episode(): Language.fromalpha2("en"), "https://subf2m.co/subtitles/breaking-bad-first-season/english/161227", "Breaking.Bad.S01E01-7.DSR-HDTV.eng", - {"title", "series", "season", "episode"}, + 7, )