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)