Merge pull request 'feature/193-process-list-of-episodes' (!1) from feature/193-process-list-of-episodes into master

Reviewed-on: #1

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_path(args.action, args.target, overrides=dict(args.overrides), season=args.season)
-        process_dir(args.target)
-    else:
-        process_file(args.target)
 
     return 0
 
 
-def process_dir(dir_path):
+def _is_media_file(path):
-    for fname in os.listdir(dir_path):
+    if os.path.isdir(path):
-        fpath = os.path.join(dir_path, fname)
+        return False
-        process_file(fpath)
+    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):
+    if os.path.isdir(path):
-        _lg.debug("Not a file: %s", fpath)
+        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
+    ep_idx = chunk_order.index("episode") + 1
-    chunk_order = chunk_order[:episode_idx] + ["episode_name"] + chunk_order[episode_idx:]
+    chunk_order = chunk_order[:ep_idx] + ["episode_name"] + chunk_order[ep_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)
 
-    if result != fname:
+    result = []
-        _lg.warning("%s -> %s", fname, 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 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):
+def _split_combined_episode(chunk_value):
-    for pat_type, pattern in PATTERNS:
+    """ 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):
+        if re.match(full_match_pat, chunk_value, flags=re.I):
-            return pat_type
+            return chunk_type
     raise RuntimeError("unhandled pattern type")
 
 

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"],
+            },
+        )

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()