diff --git a/renamer.py b/renamer.py index 3667098..9ff9d36 100755 --- a/renamer.py +++ b/renamer.py @@ -2,9 +2,11 @@ import argparse import collections +import enum import logging import os import os.path +import pprint import re import sys @@ -18,7 +20,8 @@ SEPARATORS = r"[() .!,_\[\]]" SEPARATORS_HYPHEN = r"[\-" + SEPARATORS[1:] LANGUAGES = r"(rus|eng|ukr|jap|ita|chi|kor|ger|fre|spa|pol)" PATTERNS = ( - ("episode", r"s\d{1,2}(e\d{1,2})?"), + ("season", r"s\d{1,2}"), + ("episode", r"(s\d{1,2})?e\d{1,2}"), ("year", r"(19|20)\d{2}"), ("edition", r"((theatrical|director'*s|extended|un)[-.]?cut" r"|imax[-.]edition" @@ -41,14 +44,70 @@ PATTERNS = ( ("file_extension", r"mkv|avi"), ("unknown", r".*") ) +VALID_CHUNK_TYPES = frozenset({k for k, _ in PATTERNS} | {"name", "episode_name"}) + + +# noinspection PyInterpreter +class EnumAction(argparse.Action): + """ + Argparse action for handling Enums + """ + def __init__(self, **kwargs): + # Pop off the type value + enum_type = kwargs.pop("type", None) + + # Ensure an Enum subclass is provided + if enum_type is None: + raise ValueError("type must be assigned an Enum when using EnumAction") + if not issubclass(enum_type, enum.Enum): + raise TypeError("type must be an Enum when using EnumAction") + + # Generate choices from the Enum + kwargs.setdefault("choices", tuple(e.value for e in enum_type)) + + super(EnumAction, self).__init__(**kwargs) + + self._enum = enum_type + + def __call__(self, parser, namespace, values, option_string=None): + # Convert value back into an Enum + value = self._enum(values) + setattr(namespace, self.dest, value) + + +class CliAction(enum.Enum): + parse = "parse" + rename = "rename" + _lg = logging.getLogger("spqr.movie-renamer") +def _parse_set_arg(value): + """ Parse --set TYPE=VALUE into (chunk_type, chunk_value). """ + chunk_type, sep, chunk_value = value.partition("=") + if not sep: + raise argparse.ArgumentTypeError("expected TYPE=VALUE, got %r" % value) + if chunk_type not in VALID_CHUNK_TYPES: + raise argparse.ArgumentTypeError( + "unknown chunk type %r, expected one of: %s" + % (chunk_type, ", ".join(sorted(VALID_CHUNK_TYPES))) + ) + return chunk_type, chunk_value + + def main(): parser = argparse.ArgumentParser(description="Rename media files.") - parser.add_argument("target", type=str, + parser.add_argument("action", type=CliAction, action=EnumAction, metavar="ACTION", + help="what to do with media file/directory (%(choices)s)") + parser.add_argument("target", type=str, metavar="TARGET", help="path to the media file/directory") + parser.add_argument("--set", dest="overrides", type=_parse_set_arg, action="append", + default=[], metavar="TYPE=VALUE", + help="override a parsed chunk, e.g. --set name=Foo.Bar (repeatable)") + parser.add_argument("--season", type=int, default=None, metavar="N", + help="season number; auto-numbers episodes (E01, E02, ...) " + "across media files found under TARGET") parser.add_argument("-v", "--verbose", action="store_true", default=False, help="verbose output") args = parser.parse_args() @@ -56,54 +115,103 @@ def main(): loglevel = logging.DEBUG if args.verbose else logging.INFO logging.basicConfig(level=loglevel) - if os.path.isdir(args.target): - process_dir(args.target) - else: - process_file(args.target) + process_path(args.action, args.target, overrides=dict(args.overrides), season=args.season) return 0 -def process_dir(dir_path): - for fname in os.listdir(dir_path): - fpath = os.path.join(dir_path, fname) - process_file(fpath) +def _is_media_file(path): + if os.path.isdir(path): + return False + ext = os.path.splitext(path)[1][1:] + return ext.lower() in PROCESSED_FILETYPES -def process_file(fpath): +def process_path(action: CliAction, path, overrides=None, season=None, episode=None): + overrides = overrides or {} + # process only files - if not os.path.isfile(fpath): - _lg.debug("Not a file: %s", fpath) + if os.path.isdir(path): + episode_num = 0 + for child_name in sorted(os.listdir(path)): + child_path = os.path.join(path, child_name) + child_episode = None + if season is not None and _is_media_file(child_path): + episode_num += 1 + child_episode = episode_num + process_path(action, child_path, overrides=overrides, season=season, + episode=child_episode) return + if not _is_media_file(path): + _lg.debug("Extension is not supported: %s", path) + return + + file_overrides = overrides + if season is not None: + file_overrides = dict(overrides) + file_overrides["season"] = "S%02d" % season + file_overrides["episode"] = "E%02d" % (episode if episode is not None else 1) + # split filepath to dir path, title, and extension - dir_path, fname = os.path.split(fpath) + dir_path, fname = os.path.split(path) title, ext = os.path.splitext(fname) ext = ext[1:] - if ext not in PROCESSED_FILETYPES: - _lg.debug("Extension is not supported: %s", fpath) - return parsed_title = parse_title(title) + if "name" in file_overrides and "name" not in parsed_title: + # nothing else in the title was recognized, so the whole thing + # landed in "unknown" instead of "name" -- the override replaces it + parsed_title.pop("unknown", None) + for chunk_type, chunk_value in file_overrides.items(): + parsed_title[chunk_type] = [chunk_value] - # create file name from parsed chunks + if action == CliAction.parse: + print_parsed_title(title, parsed_title) + return + + if action == CliAction.rename: + pretty_title = generate_pretty_name(parsed_title) + pretty_title += ".%s" % ext + if pretty_title != fname: + _lg.warning("%s -> %s", fname, pretty_title) + os.rename(path, os.path.join(dir_path, pretty_title)) + return + + +def print_parsed_title(title, parsed): + print(title) + pprint.pprint(parsed, indent=4) + + +def generate_pretty_name(parsed_title): + """ Create file name from parsed chunks. """ chunk_order = [k for k, _ in PATTERNS] chunk_order = ["name"] + chunk_order - episode_idx = chunk_order.index("episode") + 1 - chunk_order = chunk_order[:episode_idx] + ["episode_name"] + chunk_order[episode_idx:] - result = [] - for chunk_type in chunk_order: - if not parsed_title.get(chunk_type, []): - continue - result.append(".".join(parsed_title[chunk_type])) - result.append(ext) - result = ".".join(result) + ep_idx = chunk_order.index("episode") + 1 + chunk_order = chunk_order[:ep_idx] + ["episode_name"] + chunk_order[ep_idx:] - if result != fname: - _lg.warning("%s -> %s", fname, result) + result = [] + pending_season = None + for chunk_type in chunk_order: + chunk_values = parsed_title.get(chunk_type, []) + if not chunk_values: + continue + chunk_str = ".".join(chunk_values) + if chunk_type == "season": + pending_season = chunk_str + continue + if chunk_type == "episode" and pending_season is not None: + chunk_str = pending_season + chunk_str + pending_season = None + result.append(chunk_str) + if pending_season is not None: + result.append(pending_season) + return ".".join(result) def _get_parsed_title_dict(chunk_list, chunk_map): + """ Get {chunk_type: [chunk_value_1, ..., chunk_value_n]} dictionary. """ p_title = collections.defaultdict(list) for idx, chunk in enumerate(chunk_list): chunk_type = chunk_map[idx] @@ -112,7 +220,7 @@ def _get_parsed_title_dict(chunk_list, chunk_map): def _guess_combined(chunk_values, chunk_map): - """ Try to combine unknown chunks in pairs and parse them """ + """ Try to combine unknown chunks in pairs and parse them. """ is_changed = False p_title = _get_parsed_title_dict(chunk_values, chunk_map) if len(p_title["unknown"]) < 2: @@ -160,19 +268,20 @@ def _guess_combined(chunk_values, chunk_map): def parse_title(title): """ Split media title to components. """ + # split title by separators chunk_values = filter(None, re.split(SEPARATORS, title)) # remove non-word chunks (like single hyphens), but leave ampersands (&) chunk_values = list(filter(lambda ch: re.search(r"(\w|&)+", ch), chunk_values)) + chunk_map = [] # list of chunk_types # parse each chunk - chunk_map = [] for ch_value in chunk_values: chunk_map.append(guess_part(ch_value)) _, chunk_values, chunk_map = _guess_combined(chunk_values, chunk_map) - # # try to parse unknown chunks, replacing all hyphens in them with dots + # try to parse unknown chunks, replacing all hyphens in them with dots p_title = _get_parsed_title_dict(chunk_values, chunk_map) is_changed = False if p_title.get("unknown"): @@ -207,9 +316,12 @@ def parse_title(title): while idx < len(chunk_map) and chunk_map[idx] == "unknown": chunk_map[idx] = "name" idx += 1 - # if episode number is found, next unknown chunks are episode name - if p_title.get("episode"): - idx = chunk_map.index("episode") + 1 + # if season/episode number is found, next unknown chunks are episode name + if p_title.get("season") or p_title.get("episode"): + if p_title.get("episode"): + idx = chunk_map.index("episode") + 1 + else: + idx = chunk_map.index("season") + 1 while idx < len(chunk_map) and chunk_map[idx] == "unknown": chunk_map[idx] = "episode_name" idx += 1 @@ -226,15 +338,35 @@ def parse_title(title): continue chunk_values[idx] = chunk_value.strip("-") + # split glued SxxExx episode tokens (e.g. "S04E06") into separate + # season and episode chunks p_title = _get_parsed_title_dict(chunk_values, chunk_map) + if p_title.get("episode"): + bare_episodes = [] + for ep_value in p_title["episode"]: + if ep_value[0].lower() != "s": + bare_episodes.append(ep_value) + continue + season, episode = _split_combined_episode(ep_value) + p_title["season"].append(season) + bare_episodes.append(episode) + p_title["episode"] = bare_episodes + return dict(p_title) -def guess_part(fname_part): - for pat_type, pattern in PATTERNS: +def _split_combined_episode(chunk_value): + """ Split a glued SxxExx token like "S04E06" into ("S04", "E06"). """ + match = re.match(r"(s\d{1,2})(e\d{1,2})$", chunk_value, flags=re.I) + return match.group(1), match.group(2) + + +def guess_part(chunk_value): + """ Return chunk type for given chunk value. """ + for chunk_type, pattern in PATTERNS: full_match_pat = r"^" + pattern + r"$" - if re.match(full_match_pat, fname_part, flags=re.I): - return pat_type + if re.match(full_match_pat, chunk_value, flags=re.I): + return chunk_type raise RuntimeError("unhandled pattern type") diff --git a/tests/test_parsing.py b/tests/test_parsing.py index bdb235f..6e50e0c 100644 --- a/tests/test_parsing.py +++ b/tests/test_parsing.py @@ -10,7 +10,8 @@ class TestParserParts(unittest.TestCase): self.assertEqual( res, { - "episode": ["S04E06"], + "season": ["S04"], + "episode": ["E06"], "resolution": ["1080p"], "quality": ["WEB-DL"], "episode_name": ["Live", "Bait"], @@ -23,11 +24,13 @@ class TestParserParts(unittest.TestCase): def test_episode_number(self): title = "Vikings.S01E01.720p.BluRay.4xRus.Eng.HDCLUB" res = parse_title(title) - self.assertIn("S01E01", res.get("episode", [])) + self.assertEqual(["S01"], res.get("season")) + self.assertEqual(["E01"], res.get("episode")) self.assertEqual( res, { - "episode": ["S01E01"], + "season": ["S01"], + "episode": ["E01"], "resolution": ["720p"], "quality": ["BluRay"], "language": ["4xRus", "Eng"], @@ -59,7 +62,8 @@ class TestParserParts(unittest.TestCase): res, { "name": ["The", "Guild"], - "episode": ["s04e06"], + "season": ["s04"], + "episode": ["e06"], "episode_name": ["Weird", "Respawn", "by", "Swich"], "file_extension": ["mkv"], }, @@ -266,7 +270,8 @@ class TestCornerCases(unittest.TestCase): res, { "name": ["The", "IT", "Crowd"], - "episode": ["S01E04"], + "season": ["S01"], + "episode": ["E04"], "episode_name": ["The", "Red", "Door", "HR"], "quality": ["DVDRip"], "edition": ["HQ.Edition"], @@ -284,7 +289,8 @@ class TestCornerCases(unittest.TestCase): { "name": ["The", "Big", "Bang", "Theory"], "year": ["2019"], - "episode": ["S12E20"], + "season": ["S12"], + "episode": ["E20"], "episode_name": ["The", "Decision", "Reverberation"], "unknown": ["NTb", "EniaHD"], "resolution": ["1080p"], @@ -303,7 +309,8 @@ class TestCornerCases(unittest.TestCase): res, { "name": ["The", "Big", "Bang", "Theory"], - "episode": ["S04E06"], + "season": ["S04"], + "episode": ["E06"], "resolution": ["720p"], "quality": ["WEB-DL"], "language": ["eng", "rus"], @@ -318,8 +325,26 @@ class TestCornerCases(unittest.TestCase): self.assertEqual( res, { - "episode": ["S27E01"], + "season": ["S27"], + "episode": ["E01"], "episode_name": ["Every", "Man's", "Dream"], "file_extension": ["mkv"], }, ) + + def test_season_episode_divided(self): + title = "Lethal.Weapon.S01.E18.2016.1080p.WEBRip.Rus.Eng.HDCLUB" + res = parse_title(title) + self.assertEqual( + res, + { + "season": ["S01"], + "episode": ["E18"], + "language": ["Rus", "Eng"], + "name": ["Lethal", "Weapon"], + "quality": ["WEBRip"], + "resolution": ["1080p"], + "unknown": ["HDCLUB"], + "year": ["2016"], + }, + ) diff --git a/tests/test_process_path.py b/tests/test_process_path.py new file mode 100644 index 0000000..87f0710 --- /dev/null +++ b/tests/test_process_path.py @@ -0,0 +1,100 @@ +import os +import tempfile +import unittest + +from renamer import CliAction, _is_media_file, process_path + + +class TestIsMediaFile(unittest.TestCase): + def test_extension_is_case_insensitive(self): + self.assertTrue(_is_media_file("Show.S04E06.MKV")) + self.assertTrue(_is_media_file("Show.S04E06.mkv")) + self.assertTrue(_is_media_file("Show.S04E06.AVI")) + + def test_unsupported_extension(self): + self.assertFalse(_is_media_file("Show.S04E06.txt")) + + +class TestRenameCaseInsensitiveExtension(unittest.TestCase): + def test_uppercase_extension_is_processed(self): + with tempfile.TemporaryDirectory() as root: + original = "Show S04E06.MKV" + open(os.path.join(root, original), "w").close() + + process_path(CliAction.rename, os.path.join(root, original)) + + self.assertEqual(os.listdir(root), ["Show.S04E06.MKV"]) + + +class TestSeasonNumbering(unittest.TestCase): + def test_scoped_to_directory_not_whole_tree(self): + with tempfile.TemporaryDirectory() as root: + season_dir = os.path.join(root, "Season 1") + extras_dir = os.path.join(root, "Extras") + os.mkdir(season_dir) + os.mkdir(extras_dir) + for name in ( + "Show.(01.serija.iz.10).mkv", + "Show.(02.serija.iz.10).mkv", + ): + open(os.path.join(season_dir, name), "w").close() + open(os.path.join(extras_dir, "Blooper.reel.mkv"), "w").close() + + process_path( + CliAction.rename, root, overrides={"name": "Show"}, season=1, + ) + + self.assertEqual( + sorted(os.listdir(season_dir)), + ["Show.S01E01.mkv", "Show.S01E02.mkv"], + ) + # Extras is a separate directory: it gets its own numbering, + # starting fresh at E01, instead of continuing the season's count. + self.assertEqual(os.listdir(extras_dir), ["Show.S01E01.mkv"]) + + def test_single_file_target_defaults_to_episode_one(self): + with tempfile.TemporaryDirectory() as root: + original = "Show.(01.serija.iz.10).mkv" + open(os.path.join(root, original), "w").close() + + process_path( + CliAction.rename, os.path.join(root, original), + overrides={"name": "Show"}, season=1, + ) + + self.assertEqual(os.listdir(root), ["Show.S01E01.mkv"]) + + +class TestSetOverride(unittest.TestCase): + def test_name_override_drops_unrecognized_leftovers(self): + with tempfile.TemporaryDirectory() as root: + # nothing else in the title is recognized (no year/quality/etc), + # so it would otherwise land entirely under "unknown" + original = "Show.(01.serija.iz.10).mkv" + open(os.path.join(root, original), "w").close() + + process_path( + CliAction.rename, os.path.join(root, original), + overrides={"name": "Show"}, season=1, + ) + + self.assertEqual(os.listdir(root), ["Show.S01E01.mkv"]) + + def test_name_override_preserves_unknown_when_name_was_recognized(self): + with tempfile.TemporaryDirectory() as root: + original = "The.Walking.Dead.S04E06.1080p.WEB-DL.Rus.Eng.HDCLUB.mkv" + open(os.path.join(root, original), "w").close() + + process_path( + CliAction.rename, os.path.join(root, original), + overrides={"name": "TWD"}, + ) + + self.assertEqual( + os.listdir(root), + ["TWD.S04E06.1080p.WEB-DL.Rus.Eng.HDCLUB.mkv"], + ) + + +if __name__ == "__main__": + unittest.main()