From 76bc4179c6c05c2c998a1de7f7e1742d25c819d6 Mon Sep 17 00:00:00 2001
From: Linus <wkchan.linus@gmail.com>
Date: Sun, 14 Jul 2019 00:37:50 +1200
Subject: [PATCH] Handled header node

---
 ffmpeg/_ffmpeg.py           | 19 ++++++++++++++-----
 ffmpeg/_run.py              | 15 ++++++++++++++-
 ffmpeg/nodes.py             |  4 ++--
 ffmpeg/tests/test_ffmpeg.py |  8 +++++++-
 4 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/ffmpeg/_ffmpeg.py b/ffmpeg/_ffmpeg.py
index fa3ae79..2cbcf86 100644
--- a/ffmpeg/_ffmpeg.py
+++ b/ffmpeg/_ffmpeg.py
@@ -14,7 +14,8 @@ from .nodes import (
 )
 
 
-def input(filename, *stream, **kwargs):
+@filter_operator()
+def input(*streams_and_filename, **kwargs):
     """Input file URL (ffmpeg ``-i`` option)
 
     Any supplied kwargs are passed to ffmpeg verbatim (e.g. ``t=20``,
@@ -24,19 +25,27 @@ def input(filename, *stream, **kwargs):
 
     Official documentation: `Main options <https://ffmpeg.org/ffmpeg.html#Main-options>`__
     """
-    kwargs['filename'] = filename
+    streams_and_filename = list(streams_and_filename)
+    if 'filename' not in kwargs:
+        if not isinstance(streams_and_filename[-1], basestring):
+            raise ValueError('A filename must be provided')
+        kwargs['filename'] = streams_and_filename.pop(-1)
+    streams = streams_and_filename
+
     fmt = kwargs.pop('f', None)
     if fmt:
         if 'format' in kwargs:
             raise ValueError("Can't specify both `format` and `f` kwargs")
         kwargs['format'] = fmt
-    return InputNode(input.__name__, stream=stream, kwargs=kwargs).stream()
+    return InputNode(name=input.__name__, stream=streams, kwargs=kwargs).stream()
 
 
-def header(**kwargs):
+def header(*args, **kwargs):
     """Add extra header command-line argument(s), e.g. ``-re``.
     """
-    return HeaderNode(header.__name__, kwargs=kwargs).stream()
+    stream = None
+    print("in header: {}, {}, {}".format(repr(stream), args, kwargs))
+    return HeaderNode(name=header.__name__, args=args, kwargs=kwargs).stream()
 
 
 @output_operator()
diff --git a/ffmpeg/_run.py b/ffmpeg/_run.py
index ac62720..a6f08de 100644
--- a/ffmpeg/_run.py
+++ b/ffmpeg/_run.py
@@ -107,6 +107,17 @@ def _get_filter_arg(filter_nodes, outgoing_edge_maps, stream_name_map):
     return ';'.join(filter_specs)
 
 
+def _get_header_args(node):
+    kwargs = copy.copy(node.kwargs)
+    args = []
+    print(node.args)
+    print(node.kwargs)
+    for arg in node.args:
+        args += arg
+    args += convert_kwargs_to_cmd_line_args(kwargs)
+    return args
+
+
 def _get_global_args(node):
     return list(node.args)
 
@@ -155,13 +166,15 @@ def get_args(stream_spec, overwrite_output=False):
     # TODO: group nodes together, e.g. `-i somefile -r somerate`.
     sorted_nodes, outgoing_edge_maps = topo_sort(nodes)
     header_nodes = [node for node in sorted_nodes if isinstance(node, HeaderNode)]
+
+    print('{}'.format(repr(header_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)]
     global_nodes = [node for node in sorted_nodes if isinstance(node, GlobalNode)]
     filter_nodes = [node for node in sorted_nodes if isinstance(node, FilterNode)]
     stream_name_map = {(node, None): str(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_global_args(node) for node in header_nodes], [])
+    args += reduce(operator.add, [_get_header_args(node) for node in header_nodes], [])
     args += reduce(operator.add, [_get_input_args(node) for node in input_nodes])
     if filter_arg:
         args += ['-filter_complex', filter_arg]
diff --git a/ffmpeg/nodes.py b/ffmpeg/nodes.py
index 9019bd3..7a23814 100644
--- a/ffmpeg/nodes.py
+++ b/ffmpeg/nodes.py
@@ -198,6 +198,7 @@ class Node(KwargReprNode):
         kwargs={},
     ):
         stream_map = get_stream_map(stream_spec)
+        print('{}'.format(stream_map))
         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)
@@ -250,7 +251,7 @@ class InputNode(Node):
             incoming_stream_types={FilterableStream},
             outgoing_stream_type=FilterableStream,
             min_inputs=0,
-            max_inputs=0,
+            max_inputs=None,
             args=args,
             kwargs=kwargs,
         )
@@ -342,7 +343,6 @@ class MergeOutputsNode(Node):
         )
 
 
-# noinspection PyMethodOverriding
 class HeaderNode (Node):
 
     def __init__(self, name: str = None, args=[], kwargs={}):
diff --git a/ffmpeg/tests/test_ffmpeg.py b/ffmpeg/tests/test_ffmpeg.py
index aecf069..58c2dd2 100644
--- a/ffmpeg/tests/test_ffmpeg.py
+++ b/ffmpeg/tests/test_ffmpeg.py
@@ -136,13 +136,19 @@ def test_global_args():
 
 
 def test_header_args():
+
+    # header = ffmpeg.header('-re', '-s', thread_queue_size=512)
+
+    # print('in test: {}'.format(repr(header)))
+
     out_file = (
         ffmpeg.header(thread_queue_size='512')
         .input("input.mp4")
         .output("output.mp4")
     )
     assert out_file.get_args() == [
-        '-thread_queue_size=512',
+        '-thread_queue_size',
+        '512',
         '-i',
         'input.mp4',
         'output.mp4',