torrent & nyaa_parts: support torrent media comprised of multiple tor…

…rents

amt/server.py

@@ -278,6 +278,8 @@ def create_media_data(self, id, name, season_id=None, season_title="", dir_name=
                 lang = "dub" if match else ""
         if media_type is None:
             media_type = self.media_type
+        if self.torrent:
+            kwargs["torrent_files"] = kwargs.get("torrent_files") or []
 
         return MediaData(dict(server_id=self.id, server_alias=self.alias, id=id, dir_name=dir_name if dir_name else re.sub(r"[\W]", "", (season_title or name).replace(" ", "_")), name=name, media_type=media_type.value, media_type_name=media_type.name, progress=0, season_id=season_id, season_title=season_title, offset=offset, alt_id=alt_id, trackers={}, progress_type=progress_type if progress_type is not None else self.progress_type, tags=[], lang=lang, nextTimeStamp=0, official=self.official, version=self.version, **kwargs))
 

amt/servers/nyaa.py

@@ -3,6 +3,7 @@
 
 from ..util.media_type import MediaType
 from .torrent import GenericTorrentServer
+from difflib import SequenceMatcher
 
 
 def category_to_media_type(cat):
@@ -28,18 +29,21 @@ class Nyaa(GenericTorrentServer):
     search_url = base_url + "/?s=size&o=desc&f=0&c={}&q={}"
     torrent_url = base_url + "/download/{}.torrent"
 
-    stream_url_regex = re.compile(r"nyaa.si/view/(\w*)")
+    stream_url_regex = re.compile(domain + r"/view/(\w*)")
     media_type = MediaType.ANIME | MediaType.NOVEL | MediaType.MANGA
 
     def get_media_list(self, **kwargs):
         return self.search_for_media("", **kwargs)
 
-    def search_for_media(self, term, media_type=None, **kwargs):
-        category = media_type_to_category(media_type)
-        r = self.session_get(self.search_url.format(category, term))
+    def get_torrent_url_from_basename(self, media_data, basename):
+        return self.torrent_url.format(basename)
+
+    def search_for_media_helper(self, term, media_type=None, url=None):
+        if not url:
+            url = self.search_url.format(media_type_to_category(media_type), term)
+        r = self.session_get_cache(url)
         soup = self.soupify(BeautifulSoup, r)
         table = soup.find("table", {"class": "torrent-list"})
-        results = []
         row_num_to_media_type = {}
         for row_num, row, link in ((row_num, row, link) for row_num, row in enumerate(table.findAll("tr")) for link in row.findAll("a")) if table else []:
             if link["href"].startswith("/?c="):
@@ -49,7 +53,12 @@ def search_for_media(self, term, media_type=None, **kwargs):
                 title = link["title"]
                 label = " ".join(filter(lambda x: x, map(lambda x: x.getText().strip(), row.findAll("td", {"class": "text-center"}))))
                 media_type = row_num_to_media_type[row_num]
-                results.append(self.create_media_data(id=slug, name=title, label=label, torrent_file=self.torrent_url.format(slug), media_type=media_type))
+                yield slug, title, media_type, label
+
+    def search_for_media(self, term, media_type=None, **kwargs):
+        results = []
+        for slug, title, mediatype, label in self.search_for_media_helper(term, media_type=media_type):
+            results.append(self.create_media_data(id=slug, name=title, label=label, torrent_files=[slug], media_type=mediatype))
         return results
 
     def get_media_data_from_url(self, url):
@@ -61,4 +70,51 @@ def get_media_data_from_url(self, url):
             if link["href"].startswith("/?c="):
                 media_type = category_to_media_type(link["href"].split("=", 2)[1])
                 break
-        return self.create_media_data(id=slug, name=title, media_type=media_type, torrent_file=self.torrent_url.format(slug))
+        return self.create_media_data(id=slug, name=title, media_type=media_type, torrent_files=[slug])
+
+
+class NyaaParts(Nyaa):
+    id = "nyaa_parts"
+    alias = "nyaa"
+    search_url = Nyaa.search_url + "&o=asc"
+    media_type = MediaType.ANIME
+
+    stream_url_regex = re.compile(Nyaa.domain + r"/.*\?.*q=.+")
+
+    def group_entries(self, entries, title, media_type):
+        for e in entries.keys():
+            if len(e) == len(title):
+                s = SequenceMatcher(None, e, title)
+                seqs = s.get_matching_blocks()
+                total_same = sum([s[-1] for s in seqs])
+                if total_same >= len(title) / 2:
+                    value = s.get_matching_blocks()[0], media_type
+                    if entries[e][0] == None:
+                        entries[e] = value
+                    elif entries[e] == value:
+                        return
+        entries[title] = None, media_type
+
+    def get_all_media_data_from_url(self, url):
+        return self.search_for_media(None, url=url)
+
+    def search_for_media(self, term, media_type=None, url=None, **kwargs):
+        results = []
+        entries = {}
+        for slug, title, mediatype, _ in self.search_for_media_helper(term, media_type=media_type, url=url):
+            if "[" in title:
+                self.group_entries(entries, title, mediatype)
+        for e in entries:
+            matches, mediatype = entries[e]
+            if matches:
+                title = " ".join(e[matches[0]:matches[2]].split(" ")[:-1])
+                alt_id = hex(abs(hash(title)))[2:]
+                results.append(self.create_media_data(id=title, alt_id=alt_id, name=title, media_type=mediatype))
+        return results
+
+    def update_media_data(self, media_data, **kwargs):
+        for slug, title, mediatype, _ in self.search_for_media_helper(media_data["name"], media_type=media_data["media_type"]):
+            if title.startswith(media_data["name"]):
+                media_data["torrent_files"].append(slug)
+        media_data["torrent_files"] = list(set(media_data["torrent_files"]))
+        super().update_media_data(media_data, **kwargs)

amt/servers/torrent.py

@@ -14,27 +14,42 @@ class GenericTorrentServer(Server):
     progress_type = ProgressType.VOLUME_ONLY
     official = False
     torrent = True
+    version = 1
+
+    def upgrade_state(self, media_data):
+        return media_data["id"]
 
     @classmethod
     def get_instances(clazz, session, settings, **kwargs):
         return super().get_instances(session, settings, **kwargs) if settings.torrent_list_cmd else []
 
+    def get_torrent_url_from_basename(self, media_data, filename):  # pragma: no cover
+        return filename
+
+    def get_abs_torrent_file_path(self, media_data, filename):
+        dir_path = self.settings.get_media_dir(media_data)
+        os.makedirs(dir_path, exist_ok=True)
+        return os.path.join(dir_path, os.path.basename(filename))
+
     def list_files(self, media_data):
         dir_path = self.settings.get_media_dir(media_data)
-        return self.settings.run_cmd_and_save_output(self.settings.torrent_list_cmd, media_data=media_data, env_extra={"TORRENT_FILE": media_data["torrent_file"]}, wd=dir_path).splitlines()
+        for torrent_file in media_data["torrent_files"]:
+            abs_torrent_file_path = self.get_abs_torrent_file_path(media_data, self.get_torrent_url_from_basename(media_data, torrent_file))
+            for file in self.settings.run_cmd_and_save_output(self.settings.torrent_list_cmd, media_data=media_data, env_extra={"TORRENT_FILE": abs_torrent_file_path}, wd=dir_path).splitlines():
+                yield abs_torrent_file_path, file
+
+    def update_media_data(self, media_data, limit=None, **kwargs):
+        for torrent_file in media_data["torrent_files"][:limit]:
+            self.download_torrent_file(media_data, self.get_torrent_url_from_basename(media_data, torrent_file))
 
-    def update_media_data(self, media_data):
-        if not media_data.get("downloaded_torrent_file", False):
-            self.download_torrent_file(media_data)
-            media_data["downloaded_torrent_file"] = True
-        files = self.list_files(media_data)
+        files = list(self.list_files(media_data))
 
         numbers = set()
         duplicate_numbers = set()
-        for file in files:
+        for torrent_file, file in files:
             title = os.path.basename(file)
             n = name_parser.get_number_from_file_name(file, media_name=media_data["name"])
-            self.update_chapter_data(media_data, id=file, title=title, alt_id=title, number=n, path=file)
+            self.update_chapter_data(media_data, id=file, title=title, alt_id=title, number=n, path=file, torrent_file=torrent_file)
             if n and not media_data.chapters[file]["special"]:
                 if n in numbers:
                     duplicate_numbers.add(n)
@@ -54,7 +69,7 @@ def download_pages(self, media_data, chapter_data, **kwargs):
         dir_path = self.settings.get_media_dir(media_data)
         os.makedirs(dir_path, exist_ok=True)
         assert(os.path.exists(dir_path))
-        self.settings.run_cmd(self.settings.torrent_download_cmd, media_data=media_data, chapter_data=chapter_data, wd=dir_path, raiseException=True, env_extra={"TORRENT_FILE": media_data["torrent_file"]})
+        self.settings.run_cmd(self.settings.torrent_download_cmd, media_data=media_data, chapter_data=chapter_data, wd=dir_path, raiseException=True, env_extra={"TORRENT_FILE": chapter_data["torrent_file"]})
         return [chapter_data["id"]]
 
     def post_download(self, media_data, chapter_data, page_paths):
@@ -65,20 +80,16 @@ def post_download(self, media_data, chapter_data, page_paths):
     def get_stream_url(self, media_data, chapter_data, stream_index=0):
         return [None]
 
-    def download_torrent_file(self, media_data):
+    def download_torrent_file(self, media_data, torrent_file_url):
         """
         Downloads the raw torrent file
         """
 
-        dir_path = self.settings.get_media_dir(media_data)
-        os.makedirs(dir_path, exist_ok=True)
-        torrent_file = media_data["torrent_file"]
-        assert(torrent_file)
-        path = os.path.join(dir_path, os.path.basename(torrent_file))
+        path = self.get_abs_torrent_file_path(media_data, torrent_file_url)
         if not os.path.exists(path):
             self.logger.info("Downloading torrent file to %s", path)
-            self.save_torrent_file(torrent_file, path)
-        media_data["torrent_file"] = path
+            self.save_torrent_file(torrent_file_url, path)
+        return path
 
     def save_torrent_file(self, torrent_file, path):
         if os.path.exists(torrent_file):
@@ -124,4 +135,4 @@ def get_media_data_from_url(self, url):
                 elif key == "name":
                     title = value
 
-        return self.create_media_data(id=info_hash, name=title, torrent_file=torrent_file)
+        return self.create_media_data(id=info_hash, name=title, torrent_files=[torrent_file])

amt/settings.py

@@ -367,7 +367,7 @@ def open_viewer(self, raw_files, media_chapters):
         from shlex import quote
         for f in raw_files:
             if f is None:
-                env_extra["TORRENT_FILE"] = media_data["torrent_file"]
+                env_extra["TORRENT_FILE"] = chapter_data["torrent_file"]
                 env_extra["STREAMING"] = "1"
                 torrent = True
                 f = "-"

amt/tests/test.py

@@ -2350,9 +2350,10 @@ class ServerStreamTest(RealBaseUnitTestClass):
         ("https://mangaplus.shueisha.co.jp/titles/100020", 100020),
         ("https://mangasee123.com/manga/Mairimashita-Iruma-kun", "Mairimashita-Iruma-kun"),
         ("https://mangasee123.com/manga/Gunslinger-Girl", "Gunslinger-Girl"),
+        ("https://nyaa.si/?f=0&c=1_2&q=%5BEMBER%5D+Watashi+no+Shiawase+na+Kekkon", "[EMBER] Watashi no Shiawase na Kekkon"),
+        ("https://nyaa.si/view/1047104", "1047104"),
         ("https://nyaa.si/view/135283", "135283"),
         ("https://nyaa.si/view/269191", "269191"),
-        ("https://nyaa.si/view/1047104", "1047104"),
         ("https://viz.com/shonenjump/chapters/my-hero-academia-vigilantes", "my-hero-academia-vigilantes"),
         ("https://webtoons.com/en/drama/lookism/list?title_no=1049", 1049),
     ]