Merge pull request kkroening#188 from akolpakov/ffprobe_extra_args

kkroening · web-flow · commit 1f3ce1e2aad8 · 2019-04-18T02:05:08.000-05:00
Ability to accept extra arguments for ffmpeg.probe command (Issue kkroening#187)
diff --git a/ffmpeg/_probe.py b/ffmpeg/_probe.py
@@ -1,9 +1,10 @@
 import json
 import subprocess
 from ._run import Error
+from ._utils import convert_kwargs_to_cmd_line_args
 
 
-def probe(filename, cmd='ffprobe'):
+def probe(filename, cmd='ffprobe', **kwargs):
     """Run ffprobe on the specified file and return a JSON representation of the output.
 
     Raises:
@@ -12,7 +13,10 @@ def probe(filename, cmd='ffprobe'):
             The stderr output can be retrieved by accessing the
             ``stderr`` property of the exception.
     """
-    args = [cmd, '-show_format', '-show_streams', '-of', 'json', filename]
+    args = [cmd, '-show_format', '-show_streams', '-of', 'json']
+    args += convert_kwargs_to_cmd_line_args(kwargs)
+    args += [filename]
+
     p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
     out, err = p.communicate()
     if p.returncode != 0:
diff --git a/ffmpeg/_run.py b/ffmpeg/_run.py
@@ -1,6 +1,6 @@
 from __future__ import unicode_literals
 from .dag import get_outgoing_edges, topo_sort
-from ._utils import basestring
+from ._utils import basestring, convert_kwargs_to_cmd_line_args
 from builtins import str
 from functools import reduce
 import collections
@@ -29,16 +29,6 @@ def __init__(self, cmd, stdout, stderr):
         self.stderr = stderr
 
 
-def _convert_kwargs_to_cmd_line_args(kwargs):
-    args = []
-    for k in sorted(kwargs.keys()):
-        v = kwargs[k]
-        args.append('-{}'.format(k))
-        if v is not None:
-            args.append('{}'.format(v))
-    return args
-
-
 def _get_input_args(input_node):
     if input_node.name == input.__name__:
         kwargs = copy.copy(input_node.kwargs)
@@ -50,7 +40,7 @@ def _get_input_args(input_node):
             args += ['-f', fmt]
         if video_size:
             args += ['-video_size', '{}x{}'.format(video_size[0], video_size[1])]
-        args += _convert_kwargs_to_cmd_line_args(kwargs)
+        args += convert_kwargs_to_cmd_line_args(kwargs)
         args += ['-i', filename]
     else:
         raise ValueError('Unsupported input node: {}'.format(input_node))
@@ -136,7 +126,7 @@ def _get_output_args(node, stream_name_map):
         if not isinstance(video_size, basestring) and isinstance(video_size, collections.Iterable):
             video_size = '{}x{}'.format(video_size[0], video_size[1])
         args += ['-video_size', video_size]
-    args += _convert_kwargs_to_cmd_line_args(kwargs)
+    args += convert_kwargs_to_cmd_line_args(kwargs)
     args += [filename]
     return args
 
diff --git a/ffmpeg/_utils.py b/ffmpeg/_utils.py
@@ -78,3 +78,14 @@ def escape_chars(text, chars):
     for ch in chars:
         text = text.replace(ch, '\\' + ch)
     return text
+
+
+def convert_kwargs_to_cmd_line_args(kwargs):
+    """Helper function to build command line arguments out of dict."""
+    args = []
+    for k in sorted(kwargs.keys()):
+        v = kwargs[k]
+        args.append('-{}'.format(k))
+        if v is not None:
+            args.append('{}'.format(v))
+    return args
diff --git a/ffmpeg/tests/test_ffmpeg.py b/ffmpeg/tests/test_ffmpeg.py
@@ -650,3 +650,8 @@ def test__probe__exception():
         ffmpeg.probe(BOGUS_INPUT_FILE)
     assert str(excinfo.value) == 'ffprobe error (see stderr output for detail)'
     assert 'No such file or directory'.encode() in excinfo.value.stderr
+
+
+def test__probe__extra_args():
+    data = ffmpeg.probe(TEST_INPUT_FILE1, show_frames=None)
+    assert set(data.keys()) == {'format', 'streams', 'frames'}
