From 57fd5955e174c1b2cd68594e2f6acfab7c6cf6ba Mon Sep 17 00:00:00 2001
From: Ethan Dalool <git@voussoir.net>
Date: Wed, 19 Feb 2020 15:19:30 -0800
Subject: [PATCH] Refactor to perform both audios and subs in a single command.

---
 ffstreams.py | 55 +++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 37 insertions(+), 18 deletions(-)

diff --git a/ffstreams.py b/ffstreams.py
index f642794..d308b03 100644
--- a/ffstreams.py
+++ b/ffstreams.py
@@ -24,13 +24,13 @@ SUBTITLE_EXTENSIONS = {
 
 FFMPEG = winwhich.which('ffmpeg')
 
-def ffextractor(input_filename, prefix, search_pattern, extension_map, moveto=None):
-    input_filename = pathclass.Path(input_filename)
+def make_maps(input_file, prefix, search_pattern, extension_map, moveto=None):
+    input_file = pathclass.Path(input_file)
 
     if moveto is not None:
         moveto = pathclass.Path(moveto)
 
-    command = [FFMPEG, '-i', input_filename.absolute_path]
+    command = [FFMPEG, '-i', input_file.absolute_path]
 
     try:
         output = subprocess.check_output(command, stderr=subprocess.STDOUT)
@@ -46,8 +46,9 @@ def ffextractor(input_filename, prefix, search_pattern, extension_map, moveto=No
 
         (stream_index, codec) = (match.group(1), match.group(2))
         extension = extension_map.get(codec, extension_map['*'])
-        output_filename = input_filename.replace_extension('')
+        output_filename = input_file.replace_extension('')
         output_filename = output_filename.add_extension(f'{prefix}{stream_index}.{extension}')
+        print(f'{stream_index}, codec={codec}, ext=.{extension}')
 
         if moveto:
             output_filename = moveto.with_child(output_filename.basename)
@@ -59,36 +60,53 @@ def ffextractor(input_filename, prefix, search_pattern, extension_map, moveto=No
 
         args.append(output_filename.absolute_path)
         maps.extend(args)
+    return maps
 
-    command = [FFMPEG, '-i', input_filename.absolute_path, *maps]
-    print(command)
-    subprocess.run(command, stderr=subprocess.STDOUT)
-
-def ffaudios(input_filename, moveto=None):
-    ffextractor(
-        input_filename,
+def audio_maps(input_file, moveto=None):
+    return make_maps(
+        input_file,
         prefix='a',
         search_pattern=r'Stream #0:(\d+)[^\s]*: Audio: (\w+)',
         extension_map=AUDIO_EXTENSIONS,
         moveto=moveto,
     )
 
-def ffsubtitles(input_filename, moveto=None):
-    ffextractor(
-        input_filename,
+def subtitle_maps(input_file, moveto=None):
+    return make_maps(
+        input_file,
         prefix='s',
         search_pattern=r'Stream #0:(\d+)[^\s]*: Subtitle: (\w+)',
         extension_map=SUBTITLE_EXTENSIONS,
         moveto=moveto,
     )
 
+def ffstreams(input_file, do_audios=False, do_subtitles=False, dry=False, moveto=None):
+    maps = []
+    if do_audios:
+        maps.extend(audio_maps(input_file, moveto=moveto))
+
+    if do_subtitles:
+        maps.extend(subtitle_maps(input_file, moveto=moveto))
+
+    if not maps:
+        return
+
+    command = [FFMPEG, '-i', input_file.absolute_path, *maps]
+    print(command)
+    if not dry:
+        subprocess.run(command, stderr=subprocess.STDOUT)
+
 def ffstreams_argparse(args):
     for pattern in args.input_filename:
         for input_filename in winglob.glob(pattern):
-            if args.audios:
-                ffaudios(input_filename, moveto=args.moveto)
-            if args.subtitles:
-                ffsubtitles(input_filename, moveto=args.moveto)
+            input_file = pathclass.Path(input_filename)
+            ffstreams(
+                input_file,
+                do_audios=args.audios,
+                do_subtitles=args.subtitles,
+                dry=args.dry,
+                moveto=args.moveto,
+            )
 
 def main(argv):
     parser = argparse.ArgumentParser(description=__doc__)
@@ -97,6 +115,7 @@ def main(argv):
     parser.add_argument('--moveto', dest='moveto', default=None)
     parser.add_argument('--audio', '--audios', dest='audios', action='store_true')
     parser.add_argument('--subtitles', '--subs', dest='subtitles', action='store_true')
+    parser.add_argument('--dry', dest='dry', action='store_true')
     parser.set_defaults(func=ffstreams_argparse)
 
     args = parser.parse_args(argv)