#17: add proper handling of split operator

ffmpeg/_run.py

@@ -1,6 +1,6 @@
 from __future__ import unicode_literals
 
-from .dag import topo_sort
+from .dag import get_outgoing_edges, topo_sort
 from functools import reduce
 from past.builtins import basestring
 import copy
@@ -53,16 +53,31 @@ def _get_input_args(input_node):
     return args
 
 
-def _get_filter_spec(i, node, stream_name_map):
-    stream_name = _get_stream_name('v{}'.format(i))
-    stream_name_map[node] = stream_name
-    inputs = [stream_name_map[edge.upstream_node] for edge in node.incoming_edges]
-    filter_spec = '{}{}{}'.format(''.join(inputs), node._get_filter(), stream_name)
+def _get_filter_spec(node, outgoing_edge_map, stream_name_map):
+    incoming_edges = node.incoming_edges
+    outgoing_edges = get_outgoing_edges(node, outgoing_edge_map)
+    inputs = [stream_name_map[edge.upstream_node, edge.upstream_label] for edge in incoming_edges]
+    outputs = [stream_name_map[edge.upstream_node, edge.upstream_label] for edge in outgoing_edges]
+    filter_spec = '{}{}{}'.format(''.join(inputs), node._get_filter(), ''.join(outputs))
     return filter_spec
 
 
-def _get_filter_arg(filter_nodes, stream_name_map):
-    filter_specs = [_get_filter_spec(i, node, stream_name_map) for i, node in enumerate(filter_nodes)]
+def _allocate_filter_stream_names(filter_nodes, outgoing_edge_maps, stream_name_map):
+    stream_count = 0
+    for upstream_node in filter_nodes:
+        outgoing_edge_map = outgoing_edge_maps[upstream_node]
+        for upstream_label, downstreams in outgoing_edge_map.items():
+            if len(downstreams) > 1:
+                # TODO: automatically insert `splits` ahead of time via graph transformation.
+                raise ValueError('Encountered {} with multiple outgoing edges with same upstream label {!r}; a '
+                    '`split` filter is probably required'.format(upstream_node, upstream_label))
+            stream_name_map[upstream_node, upstream_label] = _get_stream_name('s{}'.format(stream_count))
+            stream_count += 1
+
+
+def _get_filter_arg(filter_nodes, outgoing_edge_maps, stream_name_map):
+    _allocate_filter_stream_names(filter_nodes, outgoing_edge_maps, stream_name_map)
+    filter_specs = [_get_filter_spec(node, outgoing_edge_maps[node], stream_name_map) for node in filter_nodes]
     return ';'.join(filter_specs)
 
 
@@ -78,7 +93,8 @@ def _get_output_args(node, stream_name_map):
         raise ValueError('Unsupported output node: {}'.format(node))
     args = []
     assert len(node.incoming_edges) == 1
-    stream_name = stream_name_map[node.incoming_edges[0].upstream_node]
+    edge = node.incoming_edges[0]
+    stream_name = stream_name_map[edge.upstream_node, edge.upstream_label]
     if stream_name != '[0]':
         args += ['-map', stream_name]
     kwargs = copy.copy(node.kwargs)
@@ -104,8 +120,8 @@ def get_args(stream):
         isinstance(node, GlobalNode)]
     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)]
-    stream_name_map = {node: _get_stream_name(i) for i, node in enumerate(input_nodes)}
-    filter_arg = _get_filter_arg(filter_nodes, stream_name_map)
+    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])
     if filter_arg:
         args += ['-filter_complex', filter_arg]

ffmpeg/nodes.py

@@ -39,6 +39,18 @@ class Stream(object):
         return out
 
 
+def get_stream_map(stream_spec):
+    if stream_spec is None:
+        stream_map = {}
+    elif isinstance(stream_spec, Stream):
+        stream_map = {None: stream_spec}
+    elif isinstance(stream_spec, (list, tuple)):
+        stream_map = dict(enumerate(stream_spec))
+    elif isinstance(stream_spec, dict):
+        stream_map = stream_spec
+    return stream_map
+
+
 class Node(KwargReprNode):
     """Node base"""
     @classmethod
@@ -55,18 +67,6 @@ class Node(KwargReprNode):
                 raise TypeError('Expected incoming stream(s) to be of one of the following types: {}; got {}'
                     .format(_get_types_str(incoming_stream_types), type(stream)))
 
-    @classmethod
-    def __get_stream_map(cls, stream_spec):
-        if stream_spec is None:
-            stream_map = {}
-        elif isinstance(stream_spec, Stream):
-            stream_map = {None: stream_spec}
-        elif isinstance(stream_spec, (list, tuple)):
-            stream_map = dict(enumerate(stream_spec))
-        elif isinstance(stream_spec, dict):
-            stream_map = stream_spec
-        return stream_map
-
     @classmethod
     def __get_incoming_edge_map(cls, stream_map):
         incoming_edge_map = {}
@@ -76,7 +76,7 @@ class Node(KwargReprNode):
 
     def __init__(self, stream_spec, name, incoming_stream_types, outgoing_stream_type, min_inputs, max_inputs, args,
             kwargs):
-        stream_map = self.__get_stream_map(stream_spec)
+        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)
         incoming_edge_map = self.__get_incoming_edge_map(stream_map)

ffmpeg/tests/test_ffmpeg.py

@@ -93,12 +93,19 @@ def test_get_args_simple():
 
 
 def _get_complex_filter_example():
-    in_file = ffmpeg.input(TEST_INPUT_FILE)
+    split = (ffmpeg
+        .input(TEST_INPUT_FILE)
+        .vflip()
+        .split()
+    )
+    split0 = split[0]
+    split1 = split[1]
+
     overlay_file = ffmpeg.input(TEST_OVERLAY_FILE)
     return (ffmpeg
         .concat(
-            in_file.trim(start_frame=10, end_frame=20),
-            in_file.trim(start_frame=30, end_frame=40),
+            split0.trim(start_frame=10, end_frame=20),
+            split1.trim(start_frame=30, end_frame=40),
         )
         .overlay(overlay_file.hflip())
         .drawbox(50, 50, 120, 120, color='red', thickness=5)
@@ -110,21 +117,23 @@ 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_FILE,
         '-i', TEST_OVERLAY_FILE,
         '-filter_complex',
-            '[0]trim=end_frame=20:start_frame=10[v0];' \
-            '[0]trim=end_frame=40:start_frame=30[v1];' \
-            '[v0][v1]concat=n=2[v2];' \
-            '[1]hflip[v3];' \
-            '[v2][v3]overlay=eof_action=repeat[v4];' \
-            '[v4]drawbox=50:50:120:120:red:t=5[v5]',
-        '-map', '[v5]', os.path.join(SAMPLE_DATA_DIR, 'dummy2.mp4'),
+            '[0]vflip[s0];' \
+            '[s0]split[s1][s2];' \
+            '[s1]trim=end_frame=20:start_frame=10[s3];' \
+            '[s2]trim=end_frame=40:start_frame=30[s4];' \
+            '[s3][s4]concat=n=2[s5];' \
+            '[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'),
         '-y'
     ]
 
 
+
 #def test_version():
 #    subprocess.check_call(['ffmpeg', '-version'])
 
@@ -156,8 +165,8 @@ def test_custom_filter():
     node = ffmpeg.output(node, 'dummy2.mp4')
     assert node.get_args() == [
         '-i', 'dummy.mp4',
-        '-filter_complex', '[0]custom_filter=a:b:kwarg1=c[v0]',
-        '-map', '[v0]',
+        '-filter_complex', '[0]custom_filter=a:b:kwarg1=c[s0]',
+        '-map', '[s0]',
         'dummy2.mp4'
     ]
 
@@ -170,8 +179,8 @@ def test_custom_filter_fluent():
     )
     assert node.get_args() == [
         '-i', 'dummy.mp4',
-        '-filter_complex', '[0]custom_filter=a:b:kwarg1=c[v0]',
-        '-map', '[v0]',
+        '-filter_complex', '[0]custom_filter=a:b:kwarg1=c[s0]',
+        '-map', '[s0]',
         'dummy2.mp4'
     ]
 
@@ -197,8 +206,8 @@ def test_pipe():
         '-pixel_format', 'rgb24',
         '-i', 'pipe:0',
         '-filter_complex',
-            '[0]trim=start_frame=2[v0]',
-        '-map', '[v0]',
+            '[0]trim=start_frame=2[s0]',
+        '-map', '[s0]',
         '-f', 'rawvideo',
         'pipe:1'
     ]