Python MKV reprocessor with xattrs tagging

loki/toMKV.py

@@ -0,0 +1,199 @@
+#!/usr/bin/env python3
+import sys, re, os, subprocess, pathlib, xattr, langcodes
+
+# Static config
+DEBUG        = True
+mux_vers     = 1
+mux_xattr    = 'zib.mux'
+base_path    = '/Volumes/Shared/TV/Willow'
+min_size     = 0.9
+lang_default = 'eng'
+
+# CLI config
+if len(sys.argv) > 1:
+  base_path = sys.argv[1]
+
+# Init
+root = pathlib.Path(base_path)
+re_seaep = re.compile(r'\bS(\d{2,3})E(\d{2,3})\b', re.IGNORECASE)
+re_media = re.compile(r'^(.*?)(?:\.\w\w)?(?:\.sdh)?\.(srt|mp4|m4v|mkv)$', re.IGNORECASE)
+re_subs  = re.compile(r'(?:\.(\w\w))?(?:\.(sdh))?\.srt$', re.IGNORECASE)
+if DEBUG:
+  print(f"Root: {root}")
+
+for (parent, dirs, files) in root.walk():
+  if DEBUG:
+    print(f"Processing: {parent}", file=sys.stderr)
+
+  # Reset for each directory
+  media = []
+
+  # Scan this directory for media
+  for file in files:
+    info = {
+      'name': file,
+      'path': parent.joinpath(file),
+    }
+
+    # Parse media paths
+    try:
+      m = re_media.search(info['name'])
+      if m:
+        info['base']   = m.group(1)
+        info['ext']    = m.group(2).lower()
+        info['g_base'] = info['base'].lower()
+    except BaseException as e:
+      print(f"Parsing error (media): {e}", file=sys.stderr)
+      sys.exit(-1)
+
+    # Skip non-media files
+    if 'ext' not in info:
+      if DEBUG:
+        print(f"Skipping non-media file: {info['name']}", file=sys.stderr)
+      continue
+
+    # Parse season/episode data
+    try:
+      m = re_seaep.search(info['name'])
+      if m:
+        info['season']  = int(m.group(1))
+        info['episode'] = int(m.group(2))
+        info['g_seaep'] = 'S{:02d}E{:02d}'.format(info['season'], info['episode'])
+    except BaseException as e:
+      print(f"Parsing error (seaep): {e}", file=sys.stderr)
+      sys.exit(-1)
+
+    # Parse xattrs
+    try:
+      if (
+        'ext' in info and info['ext'] == 'mkv'
+         and mux_xattr in xattr.listxattr(info['path'])
+      ):
+        info['mux'] = int.from_bytes(xattr.getxattr(info['path'], mux_xattr))
+    except BaseException as e:
+      print(f"Parsing error (xattr): {e}", file=sys.stderr)
+      sys.exit(-1)
+
+    if DEBUG:
+      print(f"Found media: {info}", file=sys.stderr)
+    media.append(info)
+
+  # Find groups
+  groups  = {}
+  for file in media:
+    if 'g_base' in file:
+      g_base = file['g_base']
+      if file['g_base'] not in groups:
+        groups[g_base] = []
+      if file not in groups[g_base]:
+        groups[g_base].append(file)
+
+  # Merge groups
+  for (g_base, group) in groups.items():
+    if len(group) == 1:
+      if 'mux' in group[0] and group[0]['mux'] >= mux_vers:
+        if DEBUG:
+          print(f"Skipping non-grouped file with current mux_vers: {group[0]['path']}", file=sys.stderr)
+        continue
+
+    if DEBUG:
+      print(f"Merging group: {group[0]['base']}", file=sys.stderr)
+
+    # Construct args and file lists
+    size    = 0
+    sources = []
+    subs    = []
+    video   = []
+    temp    = group[0]['path'].parent.joinpath(f'{group[0]['base']}-remux.mkv')
+    final   = group[0]['path'].parent.joinpath(f'{group[0]['base']}.mkv')
+    args    = [ 'mkvmerge', '--default-language', lang_default, '-o', temp ]
+    for file in group:
+
+      # Track all source files
+      size += file['path'].stat().st_size
+      sources.append(file['path'])
+
+      # Find the video source
+      if file['ext'] in [ 'mkv', 'mp4', 'm4v' ]:
+        if DEBUG:
+          print(f"Found video: {file['path'].name}", file=sys.stderr)
+        video.append(file['path'])
+
+      # Parse subtitle languages
+      else:
+        if DEBUG:
+          print(f"Found subtitle: {file['path'].name}", file=sys.stderr)
+        sub = { 'path': file['path'], 'lang': lang_default, 'sdh': False }
+        m = re_subs.search(file['path'].name)
+        if m:
+          for part in m.groups():
+            if part.lower() == 'sdh':
+              sub['sdh'] = True
+            else:
+              code = langcodes.Language.get(part.lower())
+              if code:
+                sub['lang'] = code.to_alpha3()
+              else:
+                print(f"Cound not match language ({file['path']}): {part}", file=sys.stderr)
+              if DEBUG:
+                print(f"Using language: {sub['lang']}", file=sys.stderr)
+        subs.append(sub)
+
+    # We should find exactly one video source
+    if len(video) != 1:
+      print(f"Did not find exactly 1 video: {group[0]['base']}", file=sys.stderr)
+      continue
+    args.extend(video)
+
+    # Subtitles get language and SDH metadata
+    for sub in subs:
+      args.extend([ '--language', f"0:{sub['lang']}" ])
+      if sub['sdh']:
+        args.extend([ '--hearing-impaired-flag', '0' ])
+      args.append(sub['path'])
+    if DEBUG:
+      print(f"Mux command: {args}", file=sys.stderr)
+
+    # mkvmerge
+    if temp.exists():
+      print(f"Mux output exists: {temp}", file=sys.stderr)
+      continue
+    merge = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+    if DEBUG:
+      print(f"Muxing result: {merge}", file=sys.stderr)
+
+    # Verify
+    if merge.returncode != 0 or not temp.exists() or temp.stat().st_size < (min_size * size):
+      print(f"Mux error: {merge.stdout.decode('utf-8')}", file=sys.stderr)
+      continue
+
+    # Set xattr
+    try:
+      xattr.setxattr(temp, mux_xattr, int(1).to_bytes())
+    except BaseException as e:
+      print(f"xattrs error ({temp}): {e}", file=sys.stderr)
+      continue
+
+    # Remove sources
+    try:
+      for source in sources:
+        if DEBUG:
+          print(f"Removing: {source}", file=sys.stderr)
+        source.unlink()
+    except BaseException as e:
+      print(f"Unlink error: {e}", file=sys.stderr)
+      continue
+
+    # Rename
+    err = None
+    try:
+      if DEBUG:
+        print(f"Renaming: {temp} -> {final}", file=sys.stderr)
+      temp.rename(final)
+    except BaseException as e:
+      err = e
+    if not final.exists() or final.stat().st_size < (min_size * size):
+      err = 'Missing output'
+    if err is not None:
+      print(f"Rename error ({final}): {e}", file=sys.stderr)
+      continue