Parameterize the regex

Do all name parsing during the scan

loki/probe.py

@@ -6,7 +6,7 @@ def probe(path, ffprobe_exec='ffprobe', logger=None):
   # Logging
   if not logger:
     logger = logging.getLogger(__name__)
-    logging.basicConfig(level=logging.WARN, stream=sys.stderr)
+    logger.setLevel(logging.WARN)
 
   # ffprobe
   args = [ ffprobe_exec, '-hide_banner', '-show_streams', path ]

loki/remux.py

@@ -4,15 +4,18 @@ import probe
 
 # Static config
 debug         = False
+propedit_exec = 'mkvpropedit'
 mux_xattr     = 'zib.mux'
 mux_vers      = 1
 lang_xattr    = 'zib.lang'
 lang_vers     = 1
-min_size      = 0.9
 lang_default  = 'eng'
-propedit_exec = 'mkvpropedit'
-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)
+lang_codecs   = [ 'subtitle', 'audio' ]
+min_size      = 0.9
+video_ext     = [ r'mkv', r'mp4', r'm4v' ]
+subs_ext      = [ r'srt' ]
+subs_sdh      = [ r'sdh' ]
+subs_lang     = [ r'\w\w' ]
 
 # ENV config
 if 'DEBUG' in os.environ:
@@ -20,6 +23,12 @@ if 'DEBUG' in os.environ:
 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)
+root = pathlib.Path(sys.argv[1])
+
 # Logging
 logger = logging.getLogger(__name__)
 if debug:
@@ -27,12 +36,14 @@ if debug:
 else:
   logging.basicConfig(level=logging.WARN, stream=sys.stderr)
 
-# CLI config
-if len(sys.argv) < 2:
-  print(f"Usage: {sys.argv[0]} base_directory", file=sys.stderr)
-  sys.exit(-1)
-root = pathlib.Path(sys.argv[1])
-logger.debug(f"Root: {root}")
+# Build filename-metadata parsing regex
+re_media = re.compile(
+  r'^(.*?)' +
+  r'(?:\.(' + '|'.join(subs_lang)             + r'))?' +
+  r'(?:\.(' + '|'.join(subs_sdh)              + r'))?' +
+  r'\.('    + '|'.join(video_ext + subs_ext)  + r')$'  ,
+  re.IGNORECASE
+)
 
 def read_xattr(path, attr):
   value = None
@@ -60,18 +71,26 @@ def scan(files):
 
   # Scan this directory for media
   for file in files:
-    info = {
-      'name': file,
-      'path': parent.joinpath(file),
-    }
+    info = { 'name': file, 'path': parent.joinpath(file), 'lang': None, 'sdh': False }
 
     # Parse media paths
     try:
       m = re_media.search(info['name'])
-      if m:
+      if m and len(m.groups()) == 4:
         info['base']   = m.group(1)
-        info['ext']    = m.group(2).lower()
         info['g_base'] = info['base'].lower()
+
+        try:
+          lang = langcodes.Language.get(m.group(2).lower())
+          info['lang'] = lang.to_alpha3()
+        except:
+          if m.group(2):
+            logger.warning(f"Unable to match language code: {m.group(2)}")
+
+        if m.group(3):
+          info['sdh'] = m.group(3).lower() in subs_sdh
+
+        info['ext'] = m.group(4).lower()
     except BaseException as e:
       logger.error(f"Skipping file due to name error ({info['name']}): {e}")
       continue
@@ -85,8 +104,9 @@ def scan(files):
     info['mux_vers']  = read_xattr(info['path'], mux_xattr)
     info['lang_vers'] = read_xattr(info['path'], lang_xattr)
 
-    logger.debug(f"Found media: {info}")
+    # Loop end
     media.append(info)
+    logger.debug(f"Found media: {info}")
   
   # Return
   return media
@@ -114,52 +134,30 @@ def mux(media):
         logger.debug(f"Skipping single-item group with current mux_vers: {group[0]['path']}")
         continue
 
-    # Construct args and file lists
+    # Find all sources for the group
     size    = 0
-    sources = []
     subs    = []
     videos  = []
-    final   = group[0]['path'].parent.joinpath(f'{group[0]['base']}.mkv')
-    temp    = final.with_suffix('.mkv-remux')
-    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' ]:
+      if file['ext'] in video_ext:
         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)
+        subs.append(file)
+      # Track total size for post-mux verification
+      size += file['path'].stat().st_size
 
-    # We should find exactly one video source
+    # We should find exactly one video in the group
     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'])
+    # Build args
+    final   = group[0]['path'].parent.joinpath(f'{group[0]['base']}.mkv')
+    temp    = final.with_suffix('.mkv-remux')
+    args    = [ 'mkvmerge', '--default-language', lang_default, '-o', temp ]
+    args.append(videos[0]['path'])
     
     # Subtitles get language and SDH metadata
     # Index is always 0; these files contain only one stream
@@ -168,7 +166,6 @@ def mux(media):
       if sub['sdh']:
         args.extend([ '--hearing-impaired-flag', '0' ])
       args.append(sub['path'])
-    del subs
 
     # mkvmerge
     if temp.exists():
@@ -186,11 +183,11 @@ def mux(media):
     # Set xattr
     if not write_xattr(temp, mux_xattr, mux_vers):
       continue
-    video['mux_vers'] = mux_vers
+    videos[0]['mux_vers'] = mux_vers
 
     # Remove sources
     try:
-      for source in sources:
+      for source in (videos + subs):
         logger.debug(f"Removing: {source['path']}")
         source['path'].unlink()
         media.remove(source)
@@ -203,9 +200,9 @@ def mux(media):
     try:
       logger.debug(f"Renaming: {temp.name} -> {final.name}")
       temp.rename(final)
-      video['path'] = final
-      video['name'] = final.name
-      media.append(video)
+      videos[0]['path'] = final
+      videos[0]['name'] = final.name
+      media.append(videos[0])
     except BaseException as e:
       err = e
     if not final.exists():
@@ -223,20 +220,29 @@ def lang(media):
     logger.debug(f"Lang: {file['name']}")
     
     # Skip files tha do not need lang processing
+    if not file['ext'] in video_ext:
+      logger.debug(f"Skipping non-video item: {file['path']}")
+      continue
     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)
+    streams = probe.probe(file['path'])
+    for stream in streams:
+      if not 'index' in stream:
+        logger.error(f"No index for stream ({file['name']}): {stream}")
+        continue
     
     # 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']))
+      if 'codec_type' not in stream or stream['codec_type'] not in lang_codecs:
+        logger.debug(f"Skipping missing or ignored codec_type: {file['name']}:{stream['index']}")
+        continue
+      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']}:{stream['index']}")
+        updates.append(int(stream['index']))
 
     # mkvpropedit
     if updates:
@@ -244,13 +250,13 @@ def lang(media):
       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)
+        logger.error(f"propedit error ({propedit.returncode}): {propedit.stdout.decode('utf-8')}")
+        continue
     
     # Set xattr
     if not write_xattr(file['path'], lang_xattr, lang_vers):