Rewrite torrent_parser

src/module/models/torrent.py

@@ -9,4 +9,22 @@ class TorrentInfo(BaseModel):
 class FileSet(BaseModel):
     media_path: str = Field(...)
     sc_subtitle: str | None = Field(None)
-    tc_subtitle: str | None = Field(None)
+    tc_subtitle: str | None = Field(None)
+
+
+class EpisodeFile(BaseModel):
+    group: str | None = Field(None)
+    title: str = Field(...)
+    season: int = Field(...)
+    episode: int = Field(None)
+    suffix: str = Field(..., regex=r"\.(mkv|mp4)$")
+
+
+class SubtitleFile(BaseModel):
+    group: str | None = Field(None)
+    title: str = Field(...)
+    season: int = Field(...)
+    episode: int = Field(None)
+    language: str = Field(..., regex=r"(zh|zh-tw)")
+    suffix: str = Field(..., regex=r"\.(ass|srt)$")
+

src/module/parser/analyser/torrent_parser.py

@@ -1,9 +1,14 @@
 import re
 import logging
 from dataclasses import dataclass
+import os.path as unix_path
+import ntpath as win_path
+
+from module.models.torrent import EpisodeFile, SubtitleFile
 
 logger = logging.getLogger(__name__)
 
+PLATFORM = "Unix"
 
 @dataclass
 class DownloadInfo:
@@ -40,6 +45,71 @@ def rename_init(name, folder_name, season, suffix) -> DownloadInfo:
     return DownloadInfo(name, season, suffix, file_name, folder_name)
 
 
+def split_path(torrent_path: str) -> str:
+    if PLATFORM == "Windows":
+        return win_path.split(torrent_path)[-1]
+    else:
+        return unix_path.split(torrent_path)[-1]
+
+
+def get_group(group_and_title) -> tuple[str | None, str]:
+    n = re.split(r"[\[\]()【】（）]", group_and_title)
+    while "" in n:
+        n.remove("")
+    if len(n) > 1:
+        return n[0], n[1]
+    else:
+        return None, n[0]
+
+
+def get_season_and_title(season_and_title) -> tuple[str, int]:
+    title = re.sub(r"([Ss]|Season )\d{1,3}", "", season_and_title).strip()
+    try:
+        season = re.search(r"([Ss]|Season )(\d{1,3})", season_and_title, re.I).group(2)
+    except AttributeError:
+        season = 1
+    return title, int(season)
+
+
+def get_subtitle_lang(subtitle_name: str) -> str:
+    for key, value in SUBTITLE_LANG.items():
+        for v in value:
+            if v in subtitle_name:
+                return key
+
+
+def parse_torrent(torrent_path: str, season: int | None = None, file_type: str = "media") -> EpisodeFile | SubtitleFile:
+    media_path = split_path(torrent_path)
+    for rule in RULES:
+        match_obj = re.match(rule, media_path, re.I)
+        if match_obj is not None:
+            group, title = get_group(match_obj.group(1))
+            if season is None:
+                title, season = get_season_and_title(title)
+            else:
+                title, _ = get_season_and_title(title)
+            episode = int(match_obj.group(2))
+            suffix = unix_path.splitext(torrent_path)[-1]
+            if file_type == "media":
+                return EpisodeFile(
+                    group=group,
+                    title=title,
+                    season=season,
+                    episode=episode,
+                    suffix=suffix
+                )
+            elif file_type == "subtitle":
+                language = get_subtitle_lang(media_path)
+                return SubtitleFile(
+                    group=group,
+                    title=title,
+                    season=season,
+                    language=language,
+                    episode=episode,
+                    suffix=suffix
+                )
+
+
 def rename_normal(info: DownloadInfo):
     for rule in RULES:
         match_obj = re.match(rule, info.name, re.I)
@@ -165,8 +235,6 @@ def torrent_parser(
 
 
 if __name__ == '__main__':
-    title = "海盗战记 S02E17.zh.ass"
-    folder_name = "海盗战记"
-    season = 2
-    suffix = ".ass"
-    print(torrent_parser(title, folder_name, season, suffix, method="advance"))
+    title = "[Lilith-Raws] Boku no Kokoro no Yabai Yatsu - 01 [Baha][WEB-DL][1080p][AVC AAC][CHT][MP4].mp4"
+    sub = parse_torrent(title, season=1)
+    print(sub)

src/test/test_torrent_parser.py

@@ -1,61 +1,52 @@
-from module.parser.analyser import torrent_parser
+from module.parser.analyser.torrent_parser import parse_torrent
 
 
 def test_torrent_parser():
     file_name = "[Lilith-Raws] Boku no Kokoro no Yabai Yatsu - 01 [Baha][WEB-DL][1080p][AVC AAC][CHT][MP4].mp4"
-    folder_name = "我内心的糟糕念头(2023)"
-    season = 1
-    suffix = ".mp4"
-    assert torrent_parser(file_name, folder_name, season, suffix, "pn") == "Boku no Kokoro no Yabai Yatsu S01E01.mp4"
-    assert torrent_parser(file_name, folder_name, season, suffix, "advance") == "我内心的糟糕念头(2023) S01E01.mp4"
-    assert torrent_parser(file_name, folder_name, season, suffix, "none") == "[Lilith-Raws] Boku no Kokoro no Yabai Yatsu - 01 [Baha][WEB-DL][1080p][AVC AAC][CHT][MP4].mp4"
+    bf = parse_torrent(file_name)
+    assert bf.title == "Boku no Kokoro no Yabai Yatsu"
+    assert bf.group == "Lilith-Raws"
+    assert bf.episode == 1
+    assert bf.season == 1
 
     file_name = "[Sakurato] Tonikaku Kawaii S2 [01][AVC-8bit 1080p AAC][CHS].mp4"
-    folder_name = "总之就是非常可爱(2021)"
-    season = 2
-    suffix = ".mp4"
-    assert torrent_parser(file_name, folder_name, season, suffix, "pn") == "Tonikaku Kawaii S02E01.mp4"
-    assert torrent_parser(file_name, folder_name, season, suffix, "advance") == "总之就是非常可爱(2021) S02E01.mp4"
+    bf = parse_torrent(file_name)
+    assert bf.title == "Tonikaku Kawaii"
+    assert bf.group == "Sakurato"
+    assert bf.episode == 1
+    assert bf.season == 2
 
     file_name = "[SweetSub&LoliHouse] Heavenly Delusion - 01 [WebRip 1080p HEVC-10bit AAC ASSx2].mkv"
-    folder_name = "天国大魔境(2023)"
-    season = 1
-    suffix = ".mkv"
-    assert torrent_parser(file_name, folder_name, season, suffix, "pn") == "Heavenly Delusion S01E01.mkv"
-    assert torrent_parser(file_name, folder_name, season, suffix, "advance") == "天国大魔境(2023) S01E01.mkv"
-
-    file_name = "[SBSUB][Kanojo mo Kanojo][01][GB][1080P](456E234).mp4"
-    folder_name = "女友也要有"
-    season = 1
-    suffix = ".mp4"
-    assert torrent_parser(file_name, folder_name, season, suffix, "pn") == "Kanojo mo Kanojo S01E01.mp4"
-    assert torrent_parser(file_name, folder_name, season, suffix, "advance") == "女友也要有 S01E01.mp4"
+    bf = parse_torrent(file_name)
+    assert bf.title == "Heavenly Delusion"
+    assert bf.group == "SweetSub&LoliHouse"
+    assert bf.episode == 1
+    assert bf.season == 1
 
     file_name = "[SBSUB][CONAN][1082][V2][1080P][AVC_AAC][CHS_JP](C1E4E331).mp4"
-    folder_name = "名侦探柯南(1996)"
-    season = 1
-    suffix = ".mp4"
-    assert torrent_parser(file_name, folder_name, season, suffix, "pn") == "CONAN S01E1082.mp4"
-    assert torrent_parser(file_name, folder_name, season, suffix, "advance") == "名侦探柯南(1996) S01E1082.mp4"
+    bf = parse_torrent(file_name)
+    assert bf.title == "CONAN"
+    assert bf.group == "SBSUB"
+    assert bf.episode == 1082
+    assert bf.season == 1
 
-    file_name = "海盗战记 S01E01.mp4"
-    folder_name = "海盗战记(2021)"
-    season = 1
-    suffix = ".mp4"
-    assert torrent_parser(file_name, folder_name, season, suffix, "pn") == "海盗战记 S01E01.mp4"
-    assert torrent_parser(file_name, folder_name, season, suffix, "advance") == "海盗战记(2021) S01E01.mp4"
+    file_name = "海盗战记/海盗战记 S01E01.mp4"
+    bf = parse_torrent(file_name)
+    assert bf.title == "海盗战记"
+    assert bf.episode == 1
+    assert bf.season == 1
 
     file_name = "海盗战记 S01E01.zh-tw.ass"
-    folder_name = "海盗战记(2021)"
-    season = 1
-    suffix = ".ass"
-    assert torrent_parser(file_name, folder_name, season, suffix, "subtitle_pn") == "海盗战记 S01E01.zh-tw.ass"
-    assert torrent_parser(file_name, folder_name, season, suffix, "subtitle_advance") == "海盗战记(2021) S01E01.zh-tw.ass"
+    sf = parse_torrent(file_name, file_type="subtitle")
+    assert sf.title == "海盗战记"
+    assert sf.episode == 1
+    assert sf.season == 1
+    assert sf.language == "zh-tw"
 
     file_name = "海盗战记 S01E01.SC.ass"
-    folder_name = "海盗战记(2021)"
-    season = 1
-    suffix = ".ass"
-    assert torrent_parser(file_name, folder_name, season, suffix, "subtitle_pn") == "海盗战记 S01E01.zh.ass"
-    assert torrent_parser(file_name, folder_name, season, suffix, "subtitle_advance") == "海盗战记(2021) S01E01.zh.ass"
+    sf = parse_torrent(file_name, file_type="subtitle")
+    assert sf.title == "海盗战记"
+    assert sf.season == 1
+    assert sf.episode == 1
+    assert sf.language == "zh"