Handle missing stream languages

Use a module
Take ENV config

.gitignore

@@ -0,0 +1 @@
+__pycache__/

loki/probe-static.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+import os, sys, subprocess, pathlib, logging, pprint
+
+# Static config
+DEBUG      = False
+lang_undef = 'und'
+ffprobe_exec = 'ffprobe'
+
+# Env config
+if 'DEBUG' in os.environ:
+  DEBUG = bool(os.environ['DEBUG'])
+if 'EXEC' in os.environ:
+  ffprobe_exec = pathlib.Path(os.environ['EXEC'])
+
+# CLI config
+if len(sys.argv) < 2:
+  print(f"Usage: {sys.argv[0]} input_file", file=sys.stderr)
+  sys.exit(-1)
+
+# Logging
+if DEBUG:
+  logging.basicConfig(stream=sys.stderr, level=logging.DEBUG)
+else:
+  logging.basicConfig(stream=sys.stderr, level=logging.WARN)
+logger = logging.getLogger(__name__)
+
+# Parse to array of dicts of stream data
+def ffprobe(path):
+  streams = {}
+
+  # ffprobe
+  args = [ ffprobe_exec, '-hide_banner', '-show_streams', path ]
+  logger.debug(f"ffprobe command: {args}")
+  probe = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+  logger.debug(f"ffprobe result: {probe}")
+  data = probe.stdout.decode('utf-8')
+  if probe.returncode != 0:
+    logger.error(f"ffprobe error ({probe.returncode}): {probe.stderr.decode('utf-8')}")
+    logger.debug(f"ffprobe stdout: {data}")
+    sys.exit(probe.returncode)
+
+  # Process streams
+  stream = None
+  for line in data.splitlines():
+    if line == "[STREAM]":
+      logger.debug('Stream start')
+      stream = {}
+    elif line == "[/STREAM]":
+      logger.debug('Stream end')
+      i = len(streams)
+      try:
+        i = int(stream['index'])
+      except ValueError:
+        pass
+      if i in streams:
+        logger.error(f"Inconsistent stream index: {i}")
+        logger.debug(f"Streams: {steams}\nNext stream: {stream}")
+        continue
+      streams[i] = stream
+      steam = None
+    elif stream is not None:
+      logger.debug(f"\tStream data: {line}")
+      try:
+        (key, val) = line.split('=', 1)
+      except ValueError:
+        logger.error(f"Skipping unexpected stream data: {line}")
+        continue
+      stream[key] = val
+
+  # Return the set of all parsed streams
+  return streams
+
+# Execute
+streams = ffprobe(sys.argv[1])
+if len(sys.argv) == 2:
+  pprint.pp(streams)
+else:
+  for (i, stream) in streams.items():
+    for key in sys.argv[2:]:
+      if key in stream:
+        print(f"{i}: {key}={stream[key]}")
+      else:
+        print(f"{i}: {key}=")

loki/probe.py

@@ -1,54 +1,42 @@
-#!/usr/bin/env python3
-import sys, subprocess, pprint
+import sys, subprocess, logging
 
-# Static config
-DEBUG      = True
-lang_undef = 'und'
-
-# CLI config
-if len(sys.argv) < 2:
-  print(f"Usage: {sys.argv[0]} input_file", file=sys.stderr)
-  sys.exit(-1)
-
-# Parse to array of dicts of stream data
-def ffprobe(path):
+def probe(path, ffprobe_exec='ffprobe', logger=None):
   streams = []
+  
+  # Logging
+  if not logger:
+    logger = logging.getLogger(__name__)
+    logging.basicConfig(level=logging.WARN, stream=sys.stderr)
 
   # ffprobe
-  args = [ 'ffprobe', '-show_streams', path ]
-  if DEBUG:
-    print(f"Running command: {args}", file=sys.stderr)
-  probe = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
-  if DEBUG:
-    print(f"ffprobe result: {probe}", file=sys.stderr)
+  args = [ ffprobe_exec, '-hide_banner', '-show_streams', path ]
+  logger.debug(f"ffprobe command: {args}")
+  probe = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+  logger.debug(f"ffprobe result: {probe}")
   data = probe.stdout.decode('utf-8')
-  if probe.returncode !=0:
-    print(f"ffprobe error ({probe.returncode}): {data}", file=sys.stderr)
+  if probe.returncode != 0:
+    logger.error(f"ffprobe error ({probe.returncode}): {probe.stderr.decode('utf-8')}")
+    logger.debug(f"ffprobe stdout: {data}")
     sys.exit(probe.returncode)
 
-  # Process streams
+  # Parse streams
   stream = None
   for line in data.splitlines():
     if line == "[STREAM]":
-      if DEBUG:
-        print('Stream start', file=sys.stderr)
+      logger.debug('Stream start')
       stream = {}
     elif line == "[/STREAM]":
-      if DEBUG:
-        print('Stream end', file=sys.stderr)
+      logger.debug('Stream end')
       streams.append(stream)
       steam = None
     elif stream is not None:
-      if DEBUG:
-        print(f"\tStream data: {line}", file=sys.stderr)
+      logger.debug(f"\tStream data: {line}")
       try:
         (key, val) = line.split('=', 1)
       except ValueError:
-        print(f"Skipping unexpected stream data: {line}", file=sys.stderr)
+        logger.error(f"Skipping unexpected stream data: {line}")
         continue
       stream[key] = val
 
+  # Return
   return streams
-
-# Execute
-pprint.pp(ffprobe(sys.argv[1]))

loki/remux.py

@@ -0,0 +1,267 @@
+#!/usr/bin/env python3
+import sys, re, os, subprocess, pathlib, xattr, langcodes, logging
+import probe
+
+# Static config
+debug         = False
+mux_xattr     = 'zib.mux'
+mux_vers      = 1
+lang_xattr    = 'zib.lang'
+lang_vers     = 1
+min_size      = 0.9
+lang_default  = 'eng'
+propedit_exec = 'mkvpropedit'
+
+# ENV config
+if 'DEBUG' in os.environ:
+  debug = bool(os.environ['DEBUG'])
+if 'PROPEDIT' in os.environ:
+  propedit_exec = Pathlib.path(os.environ['PROPEDIT'])
+
+# CLI config
+if len(sys.argv) < 2:
+  print(f"Usage: {sys.argv[0]} base_directory", file=sys.stderr)
+  sys.exit(-1)
+
+# Logging
+logger = logging.getLogger(__name__)
+if debug:
+  logging.basicConfig(level=logging.DEBUG, stream=sys.stderr)
+else:
+  logging.basicConfig(level=logging.WARN, stream=sys.stderr)
+
+# Root path
+root = pathlib.Path(sys.argv[1])
+logger.debug(f"Root: {root}")
+
+def read_xattr(path, attr):
+  value = None
+  try:
+    if attr in xattr.listxattr(path):
+      value = int.from_bytes(xattr.getxattr(path, attr))
+      logger.debug(f"got xattr: {attr} => {value}")
+  except BaseException as e:
+    logger.error(f"xattr error ({attr}): {e}")
+  return value
+
+def write_xattr(path, attr, value):
+  error = True
+  try:
+    xattr.setxattr(path, attr, value.to_bytes())
+    error = False
+    logger.debug(f"set xattr: {attr} => {value}")
+  except BaseException as e:
+    logger.error(f"xattr error ({attr}): {e}")
+  return not error
+
+def scan(files):
+  logger.debug(f"Scanning {len(files)} files")
+  media = []
+  re_subs  = re.compile(r'(?:\.(\w\w))?(?:\.(sdh))?\.srt$', re.IGNORECASE)
+  re_media = re.compile(r'^(.*?)(?:\.\w\w)?(?:\.sdh)?\.((?:srt)|(?:mp4|m4v|mkv))$', re.IGNORECASE)
+
+  # 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:
+      logger.error(f"Skipping file due to name error ({info['name']}): {e}")
+      continue
+
+    # Skip non-media files
+    if 'ext' not in info:
+      logger.debug(f"Skipping non-media file: {info['name']}")
+      continue
+
+    # Read xattrs
+    info['mux_vers']  = read_xattr(info['path'], mux_xattr)
+    info['lang_vers'] = read_xattr(info['path'], lang_xattr)
+
+    logger.debug(f"Found media: {info}")
+    media.append(info)
+  
+  # Return
+  return media
+
+def mux(media):
+  logger.debug(f"Muxing {len(media)} files")
+  
+  # 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():
+    logger.debug(f"Merging group: {group[0]['base']}")
+    
+    # Skip groups that do not need muxing
+    if len(group) == 1:
+      if 'mux_vers' in group[0] and group[0]['mux_vers'] and group[0]['mux_vers'] >= mux_vers:
+        logger.debug(f"Skipping single-item group with current mux_vers: {group[0]['path']}")
+        continue
+
+    # Construct args and file lists
+    size    = 0
+    sources = []
+    subs    = []
+    videos  = []
+    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)
+
+      # Find the video source
+      if file['ext'] in [ 'mkv', 'mp4', 'm4v' ]:
+        logger.debug(f"Found video: {file['path'].name}")
+        videos.append(file)
+
+      # Parse subtitle languages
+      else:
+        logger.debug(f"Found subtitle: {file['path'].name}")
+        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:
+                logger.error(f"Cound not match language ({file['path']}): {part}")
+              logger.debug(f"Using language: {sub['lang']}")
+        subs.append(sub)
+
+    # We should find exactly one video source
+    if len(videos) != 1:
+      logger.error(f"Did not find exactly 1 video: {group[0]['base']}")
+      continue
+    video = videos[0]
+    del videos
+    
+    # Video input file
+    args.append(video['path'])
+    
+    # Subtitles get language and SDH metadata
+    # Index is always 0; these files contain only one stream
+    for sub in subs:
+      args.extend([ '--language', f"0:{sub['lang']}" ])
+      if sub['sdh']:
+        args.extend([ '--hearing-impaired-flag', '0' ])
+      args.append(sub['path'])
+    del subs
+
+    # mkvmerge
+    if temp.exists():
+      logger.error(f"Skipping because mux output exists: {temp}")
+      continue
+    logger.debug(f"Mux command: {args}")
+    merge = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+    logger.debug(f"Mux result: {merge}")
+
+    # Verify
+    if merge.returncode != 0 or not temp.exists() or temp.stat().st_size < (min_size * size):
+      logger.error(f"Mux error: {merge.stdout.decode('utf-8')}")
+      continue
+
+    # Set xattr
+    if not write_xattr(temp, mux_xattr, mux_vers):
+      continue
+    video['mux_vers'] = mux_vers
+
+    # Remove sources
+    try:
+      for source in sources:
+        logger.debug(f"Removing: {source['path']}")
+        source['path'].unlink()
+        media.remove(source)
+    except BaseException as e:
+      logger.error(f"Unlink error: {e}")
+      continue
+
+    # Rename
+    err = None
+    try:
+      logger.debug(f"Renaming: {temp.name} -> {final.name}")
+      temp.rename(final)
+      video['path'] = final
+      video['name'] = final.name
+      media.append(video)
+    except BaseException as e:
+      err = e
+    if not final.exists():
+      err = 'Output file missing'
+    elif final.stat().st_size < (min_size * size):
+      err = 'Output file too small'
+    if err is not None:
+      logger.error(f"Rename error ({final}): {e}")
+      continue
+
+def lang(media):
+  logger.debug(f"Language processing {len(media)} files")
+  
+  for file in media:
+    logger.debug(f"Lang: {file['name']}")
+
+    # Skip files tha do not need lang processing
+    if 'lang_vers' in file and file['lang_vers'] and file['lang_vers'] >= lang_vers:
+      logger.debug(f"Skipping item with current lang_vers: {file['path']}")
+      continue
+    
+    # Probe streams
+    streams = probe.probe(file['path'], logger=logger)
+    
+    # Find missing and undefined languages for subtitle and audio streams
+    updates = []
+    for stream in streams:
+      if all(name in stream for name in [ 'codec_type', 'index' ]) and stream['codec_type'] in [ 'subtitle', 'audio' ]:
+        if 'TAG:language' not in stream or stream['TAG:language'] in [ None, 'und' ]:
+          logger.debug(f"Missing or undefined language in {file['name']}:{stream['codec_type'][0]}:{stream['index']}")
+          updates.append(int(stream['index']))
+
+    # mkvpropedit
+    if updates:
+      args = [ propedit_exec, file['path'] ]
+      for index in updates:
+        # mkvpropedit starts counting at 1
+        args.extend([ '--edit', f"track:{index + 1}", '--set', f"language={lang_default}" ])
+    
+      logger.debug(f"propedit command: {args}")
+      propedit = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+      logger.debug(f"propedit result: {propedit}")
+      if propedit.returncode != 0:
+        logger.error(f"ffprobe error ({propedit.returncode}): {propedit.stdout.decode('utf-8')}")
+        sys.exit(propedit.returncode)
+    
+    # Set xattr
+    if not write_xattr(file['path'], lang_xattr, lang_vers):
+      continue
+    file['lang_vers'] = lang_vers
+    
+# Execute
+for (parent, dirs, files) in root.walk():
+  logger.debug(f"Walking: {parent}")
+  media = scan(files)
+  mux(media)
+  lang(media)