From 5d78a2595d2c11236726e456e4887448751475c7 Mon Sep 17 00:00:00 2001 From: Karl Kroening Date: Sun, 9 Jul 2017 15:50:51 -0600 Subject: [PATCH] #17: fix `merge_outputs`; allow `stream_spec` in `get_args`+`run` --- .gitignore | 2 +- ffmpeg/_run.py | 27 +++--- ffmpeg/nodes.py | 24 ++++-- .../tests/sample_data/{dummy.mp4 => in1.mp4} | Bin ffmpeg/tests/test_ffmpeg.py | 80 +++++++++++++----- 5 files changed, 96 insertions(+), 37 deletions(-) rename ffmpeg/tests/sample_data/{dummy.mp4 => in1.mp4} (100%) diff --git a/.gitignore b/.gitignore index f0d6df7..3179ebc 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,6 @@ .eggs .tox/ dist/ -ffmpeg/tests/sample_data/dummy2.mp4 +ffmpeg/tests/sample_data/out*.mp4 ffmpeg_python.egg-info/ venv* diff --git a/ffmpeg/_run.py b/ffmpeg/_run.py index eb00c90..83117e4 100644 --- a/ffmpeg/_run.py +++ b/ffmpeg/_run.py @@ -13,11 +13,12 @@ from ._ffmpeg import ( overwrite_output, ) from .nodes import ( + get_stream_spec_nodes, + FilterNode, GlobalNode, InputNode, OutputNode, output_operator, - Stream, ) @@ -108,18 +109,16 @@ def _get_output_args(node, stream_name_map): @output_operator() -def get_args(stream): +def get_args(stream_spec, overwrite_output=False): """Get command-line arguments for ffmpeg.""" - if not isinstance(stream, Stream): - raise TypeError('Expected Stream; got {}'.format(type(stream))) + nodes = get_stream_spec_nodes(stream_spec) args = [] # TODO: group nodes together, e.g. `-i somefile -r somerate`. - sorted_nodes, outgoing_edge_maps = topo_sort([stream.node]) + sorted_nodes, outgoing_edge_maps = topo_sort(nodes) input_nodes = [node for node in sorted_nodes if isinstance(node, InputNode)] - output_nodes = [node for node in sorted_nodes if isinstance(node, OutputNode) and not - isinstance(node, GlobalNode)] + output_nodes = [node for node in sorted_nodes if isinstance(node, OutputNode)] global_nodes = [node for node in sorted_nodes if isinstance(node, GlobalNode)] - filter_nodes = [node for node in sorted_nodes if node not in (input_nodes + output_nodes + global_nodes)] + filter_nodes = [node for node in sorted_nodes if isinstance(node, FilterNode)] stream_name_map = {(node, None): _get_stream_name(i) for i, node in enumerate(input_nodes)} filter_arg = _get_filter_arg(filter_nodes, outgoing_edge_maps, stream_name_map) args += reduce(operator.add, [_get_input_args(node) for node in input_nodes]) @@ -127,17 +126,23 @@ def get_args(stream): args += ['-filter_complex', filter_arg] args += reduce(operator.add, [_get_output_args(node, stream_name_map) for node in output_nodes]) args += reduce(operator.add, [_get_global_args(node) for node in global_nodes], []) + if overwrite_output: + args += ['-y'] return args @output_operator() -def run(node, cmd='ffmpeg'): - """Run ffmpeg on node graph.""" +def run(stream_spec, cmd='ffmpeg', **kwargs): + """Run ffmpeg on node graph. + + Args: + **kwargs: keyword-arguments passed to ``get_args()`` (e.g. ``overwrite_output=True``). + """ if isinstance(cmd, basestring): cmd = [cmd] elif type(cmd) != list: cmd = list(cmd) - args = cmd + node.get_args() + args = cmd + get_args(stream_spec, **kwargs) _subprocess.check_call(args) diff --git a/ffmpeg/nodes.py b/ffmpeg/nodes.py index c2b525e..11b8f85 100644 --- a/ffmpeg/nodes.py +++ b/ffmpeg/nodes.py @@ -52,6 +52,20 @@ def get_stream_map(stream_spec): return stream_map +def get_stream_map_nodes(stream_map): + nodes = [] + for stream in stream_map.values(): + if not isinstance(stream, Stream): + raise TypeError('Expected Stream; got {}'.format(type(stream))) + nodes.append(stream.node) + return nodes + + +def get_stream_spec_nodes(stream_spec): + stream_map = get_stream_map(stream_spec) + return get_stream_map_nodes(stream_map) + + class Node(KwargReprNode): """Node base""" @classmethod @@ -75,8 +89,8 @@ class Node(KwargReprNode): incoming_edge_map[downstream_label] = (upstream.node, upstream.label) return incoming_edge_map - def __init__(self, stream_spec, name, incoming_stream_types, outgoing_stream_type, min_inputs, max_inputs, args, - kwargs): + def __init__(self, stream_spec, name, incoming_stream_types, outgoing_stream_type, min_inputs, max_inputs, args=[], + kwargs={}): stream_map = get_stream_map(stream_spec) self.__check_input_len(stream_map, min_inputs, max_inputs) self.__check_input_types(stream_map, incoming_stream_types) @@ -164,13 +178,13 @@ class OutputNode(Node): class OutputStream(Stream): def __init__(self, upstream_node, upstream_label): - super(OutputStream, self).__init__(upstream_node, upstream_label, {OutputNode, GlobalNode}) + super(OutputStream, self).__init__(upstream_node, upstream_label, {OutputNode, GlobalNode, MergeOutputsNode}) class MergeOutputsNode(Node): - def __init__(self, stream, name): + def __init__(self, streams, name): super(MergeOutputsNode, self).__init__( - stream_spec=None, + stream_spec=streams, name=name, incoming_stream_types={OutputStream}, outgoing_stream_type=OutputStream, diff --git a/ffmpeg/tests/sample_data/dummy.mp4 b/ffmpeg/tests/sample_data/in1.mp4 similarity index 100% rename from ffmpeg/tests/sample_data/dummy.mp4 rename to ffmpeg/tests/sample_data/in1.mp4 diff --git a/ffmpeg/tests/test_ffmpeg.py b/ffmpeg/tests/test_ffmpeg.py index 32b8be5..47f8f72 100644 --- a/ffmpeg/tests/test_ffmpeg.py +++ b/ffmpeg/tests/test_ffmpeg.py @@ -8,9 +8,10 @@ import random TEST_DIR = os.path.dirname(__file__) SAMPLE_DATA_DIR = os.path.join(TEST_DIR, 'sample_data') -TEST_INPUT_FILE = os.path.join(SAMPLE_DATA_DIR, 'dummy.mp4') +TEST_INPUT_FILE1 = os.path.join(SAMPLE_DATA_DIR, 'in1.mp4') TEST_OVERLAY_FILE = os.path.join(SAMPLE_DATA_DIR, 'overlay.png') -TEST_OUTPUT_FILE = os.path.join(SAMPLE_DATA_DIR, 'dummy2.mp4') +TEST_OUTPUT_FILE1 = os.path.join(SAMPLE_DATA_DIR, 'out1.mp4') +TEST_OUTPUT_FILE2 = os.path.join(SAMPLE_DATA_DIR, 'out2.mp4') subprocess.check_call(['ffmpeg', '-version']) @@ -94,7 +95,7 @@ def test_get_args_simple(): def _get_complex_filter_example(): split = (ffmpeg - .input(TEST_INPUT_FILE) + .input(TEST_INPUT_FILE1) .vflip() .split() ) @@ -109,7 +110,7 @@ def _get_complex_filter_example(): ) .overlay(overlay_file.hflip()) .drawbox(50, 50, 120, 120, color='red', thickness=5) - .output(TEST_OUTPUT_FILE) + .output(TEST_OUTPUT_FILE1) .overwrite_output() ) @@ -117,7 +118,7 @@ def _get_complex_filter_example(): def test_get_args_complex_filter(): out = _get_complex_filter_example() args = ffmpeg.get_args(out) - assert args == ['-i', TEST_INPUT_FILE, + assert args == ['-i', TEST_INPUT_FILE1, '-i', TEST_OVERLAY_FILE, '-filter_complex', '[0]vflip[s0];' \ @@ -128,7 +129,7 @@ def test_get_args_complex_filter(): '[1]hflip[s6];' \ '[s5][s6]overlay=eof_action=repeat[s7];' \ '[s7]drawbox=50:50:120:120:red:t=5[s8]', - '-map', '[s8]', os.path.join(SAMPLE_DATA_DIR, 'dummy2.mp4'), + '-map', '[s8]', TEST_OUTPUT_FILE1, '-y' ] @@ -139,31 +140,38 @@ def test_get_args_complex_filter(): def test_run(): - node = _get_complex_filter_example() - ffmpeg.run(node) + stream = _get_complex_filter_example() + ffmpeg.run(stream) + + +def test_run_multi_output(): + in_ = ffmpeg.input(TEST_INPUT_FILE1) + out1 = in_.output(TEST_OUTPUT_FILE1) + out2 = in_.output(TEST_OUTPUT_FILE2) + ffmpeg.run([out1, out2], overwrite_output=True) def test_run_dummy_cmd(): - node = _get_complex_filter_example() - ffmpeg.run(node, cmd='true') + stream = _get_complex_filter_example() + ffmpeg.run(stream, cmd='true') def test_run_dummy_cmd_list(): - node = _get_complex_filter_example() - ffmpeg.run(node, cmd=['true', 'ignored']) + stream = _get_complex_filter_example() + ffmpeg.run(stream, cmd=['true', 'ignored']) def test_run_failing_cmd(): - node = _get_complex_filter_example() + stream = _get_complex_filter_example() with pytest.raises(subprocess.CalledProcessError): - ffmpeg.run(node, cmd='false') + ffmpeg.run(stream, cmd='false') def test_custom_filter(): - node = ffmpeg.input('dummy.mp4') - node = ffmpeg.filter_(node, 'custom_filter', 'a', 'b', kwarg1='c') - node = ffmpeg.output(node, 'dummy2.mp4') - assert node.get_args() == [ + stream = ffmpeg.input('dummy.mp4') + stream = ffmpeg.filter_(stream, 'custom_filter', 'a', 'b', kwarg1='c') + stream = ffmpeg.output(stream, 'dummy2.mp4') + assert stream.get_args() == [ '-i', 'dummy.mp4', '-filter_complex', '[0]custom_filter=a:b:kwarg1=c[s0]', '-map', '[s0]', @@ -172,12 +180,12 @@ def test_custom_filter(): def test_custom_filter_fluent(): - node = (ffmpeg + stream = (ffmpeg .input('dummy.mp4') .filter_('custom_filter', 'a', 'b', kwarg1='c') .output('dummy2.mp4') ) - assert node.get_args() == [ + assert stream.get_args() == [ '-i', 'dummy.mp4', '-filter_complex', '[0]custom_filter=a:b:kwarg1=c[s0]', '-map', '[s0]', @@ -185,6 +193,38 @@ def test_custom_filter_fluent(): ] +def test_merge_outputs(): + in_ = ffmpeg.input('in.mp4') + out1 = in_.output('out1.mp4') + out2 = in_.output('out2.mp4') + assert ffmpeg.merge_outputs(out1, out2).get_args() == [ + '-i', 'in.mp4', 'out1.mp4', 'out2.mp4' + ] + assert ffmpeg.get_args([out1, out2]) == [ + '-i', 'in.mp4', 'out2.mp4', 'out1.mp4' + ] + + +def test_multi_passthrough(): + out1 = ffmpeg.input('in1.mp4').output('out1.mp4') + out2 = ffmpeg.input('in2.mp4').output('out2.mp4') + out = ffmpeg.merge_outputs(out1, out2) + assert ffmpeg.get_args(out) == [ + '-i', 'in1.mp4', + '-i', 'in2.mp4', + 'out1.mp4', + '-map', '[1]', # FIXME: this should not be here (see #23) + 'out2.mp4' + ] + assert ffmpeg.get_args([out1, out2]) == [ + '-i', 'in2.mp4', + '-i', 'in1.mp4', + 'out2.mp4', + '-map', '[1]', # FIXME: this should not be here (see #23) + 'out1.mp4' + ] + + def test_pipe(): width = 32 height = 32