Subf2m provider: add season pack support

1 files changed, 59 insertions, 22 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)