diff --git a/src/module/models/torrent.py b/src/module/models/torrent.py index 59ae7bcb..83dbbf9b 100644 --- a/src/module/models/torrent.py +++ b/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) \ No newline at end of file + 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)$") + diff --git a/src/module/parser/analyser/torrent_parser.py b/src/module/parser/analyser/torrent_parser.py index e6a05165..10c8e95e 100644 --- a/src/module/parser/analyser/torrent_parser.py +++ b/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")) \ No newline at end of file + 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) diff --git a/src/test/test_torrent_parser.py b/src/test/test_torrent_parser.py index 5252e570..77fbee81 100644 --- a/src/test/test_torrent_parser.py +++ b/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"