From a5db66913903b06dff3b476da82365dbf2122b26 Mon Sep 17 00:00:00 2001
From: Ethan Dalool <git@voussoir.net>
Date: Wed, 22 Jan 2020 15:33:48 -0800
Subject: [PATCH] Rename ffaudios.py -> ffstreams.py, generalize for subs too.

---
 ffaudios.py => ffstreams.py | 48 +++++++++++++++++++++++++++++--------
 1 file changed, 38 insertions(+), 10 deletions(-)
 rename ffaudios.py => ffstreams.py (54%)

diff --git a/ffaudios.py b/ffstreams.py
similarity index 54%
rename from ffaudios.py
rename to ffstreams.py
index 739cf7b..64f920f 100644
--- a/ffaudios.py
+++ b/ffstreams.py
@@ -7,19 +7,26 @@ from voussoirkit import pathclass
 from voussoirkit import winglob
 from voussoirkit import winwhich
 
-EXTENSIONS = {
+AUDIO_EXTENSIONS = {
     'aac': 'm4a',
     'mp3': 'mp3',
     'opus': 'opus',
+    '*': 'mka',
 }
 
-ffmpeg = winwhich.which('ffmpeg')
+SUBTITLE_EXTENSIONS = {
+    'ass': 'ass',
+    'subrip': 'srt',
+    '*': 'mkv',
+}
 
-def ffaudios(input_filename, moveto=None):
+FFMPEG = winwhich.which('ffmpeg')
+
+def ffextractor(input_filename, prefix, search_pattern, extension_map, moveto=None):
     input_filename = pathclass.Path(input_filename)
     if moveto is not None:
         moveto = pathclass.Path(moveto)
-    command = [ffmpeg, '-i', input_filename.absolute_path]
+    command = [FFMPEG, '-i', input_filename.absolute_path]
 
     try:
         output = subprocess.check_output(command, stderr=subprocess.STDOUT)
@@ -29,32 +36,53 @@ def ffaudios(input_filename, moveto=None):
     output = output.decode('utf-8')
     maps = []
     for line in output.splitlines():
-        match = re.search(r'Stream #0:(\d+)[^\s]*: Audio: (\w+)', line)
+        match = re.search(search_pattern, line)
         if match is None:
             continue
         (stream_index, codec) = (match.group(1), match.group(2))
-        extension = EXTENSIONS.get(codec, 'mka')
-        output_filename = input_filename.replace_extension('').add_extension(f'{stream_index}.{extension}')
+        extension = extension_map.get(codec, extension_map['*'])
+        output_filename = input_filename.replace_extension('').add_extension(f'{prefix}{stream_index}.{extension}')
         if moveto:
             output_filename = moveto.with_child(output_filename.basename)
         maps.extend(['-map', f'0:{stream_index}', '-c', 'copy', output_filename.absolute_path])
-        # print(output_filename)
 
-    command = [ffmpeg, '-i', input_filename.absolute_path, *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,
+        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,
+        prefix='s',
+        search_pattern=r'Stream #0:(\d+)[^\s]*: Subtitle: (\w+)',
+        extension_map=SUBTITLE_EXTENSIONS,
+        moveto=moveto,
+    )
 
 def ffaudios_argparse(args):
     for pattern in args.input_filename:
         for input_filename in winglob.glob(pattern):
-            ffaudios(input_filename, moveto=args.moveto)
+            if args.audios:
+                ffaudios(input_filename, moveto=args.moveto)
+            if args.subtitles:
+                ffsubtitles(input_filename, moveto=args.moveto)
 
 def main(argv):
     parser = argparse.ArgumentParser(description=__doc__)
 
     parser.add_argument('input_filename', nargs='+')
     parser.add_argument('--moveto', dest='moveto', default=None)
+    parser.add_argument('--audio', '--audios', dest='audios', action='store_true')
+    parser.add_argument('--subtitles', dest='subtitles', action='store_true')
     parser.set_defaults(func=ffaudios_argparse)
 
     args = parser.parse_args(argv)