From 3f65898ac26b14bf0252de7c8e19cf6283a4c7f8 Mon Sep 17 00:00:00 2001 From: Linus Date: Sat, 13 Jul 2019 04:44:28 +1200 Subject: [PATCH] add header Node --- .gitignore | 1 + .vscode/settings.json | 5 +++++ ffmpeg/_ffmpeg.py | 13 ++++++++++--- ffmpeg/_run.py | 7 +++++-- ffmpeg/nodes.py | 25 +++++++++++++++++++++---- ffmpeg/tests/test_ffmpeg.py | 16 +++++++++++++++- requirements.txt | 1 + 7 files changed, 58 insertions(+), 10 deletions(-) create mode 100644 .vscode/settings.json diff --git a/.gitignore b/.gitignore index 780f20e..866cf60 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ ffmpeg/tests/sample_data/out*.mp4 ffmpeg_python.egg-info/ venv* build/ +*pyc \ No newline at end of file diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..b5f2aa4 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,5 @@ +{ + "python.linting.pylintEnabled": false, + "python.linting.pep8Enabled": true, + "python.linting.enabled": true +} \ No newline at end of file diff --git a/ffmpeg/_ffmpeg.py b/ffmpeg/_ffmpeg.py index 31e2b90..fa3ae79 100644 --- a/ffmpeg/_ffmpeg.py +++ b/ffmpeg/_ffmpeg.py @@ -6,6 +6,7 @@ from ._utils import basestring from .nodes import ( filter_operator, GlobalNode, + HeaderNode, InputNode, MergeOutputsNode, OutputNode, @@ -13,7 +14,7 @@ from .nodes import ( ) -def input(filename, **kwargs): +def input(filename, *stream, **kwargs): """Input file URL (ffmpeg ``-i`` option) Any supplied kwargs are passed to ffmpeg verbatim (e.g. ``t=20``, @@ -29,7 +30,13 @@ def input(filename, **kwargs): if 'format' in kwargs: raise ValueError("Can't specify both `format` and `f` kwargs") kwargs['format'] = fmt - return InputNode(input.__name__, kwargs=kwargs).stream() + return InputNode(input.__name__, stream=stream, kwargs=kwargs).stream() + + +def header(**kwargs): + """Add extra header command-line argument(s), e.g. ``-re``. + """ + return HeaderNode(header.__name__, kwargs=kwargs).stream() @output_operator() @@ -94,4 +101,4 @@ def output(*streams_and_filename, **kwargs): return OutputNode(streams, output.__name__, kwargs=kwargs).stream() -__all__ = ['input', 'merge_outputs', 'output', 'overwrite_output'] +__all__ = ['header', 'input', 'merge_outputs', 'output', 'overwrite_output'] diff --git a/ffmpeg/_run.py b/ffmpeg/_run.py index afc504d..ac62720 100644 --- a/ffmpeg/_run.py +++ b/ffmpeg/_run.py @@ -13,6 +13,7 @@ from .nodes import ( get_stream_spec_nodes, FilterNode, GlobalNode, + HeaderNode, InputNode, OutputNode, output_operator, @@ -53,8 +54,8 @@ def _format_input_stream_name(stream_name_map, edge, is_final_arg=False): else: suffix = ':{}'.format(edge.upstream_selector) if is_final_arg and isinstance(edge.upstream_node, InputNode): - ## Special case: `-map` args should not have brackets for input - ## nodes. + # Special case: `-map` args should not have brackets for input + # nodes. fmt = '{}{}' else: fmt = '[{}{}]' @@ -153,12 +154,14 @@ def get_args(stream_spec, overwrite_output=False): args = [] # 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)] 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_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 cacab8e..9019bd3 100644 --- a/ffmpeg/nodes.py +++ b/ffmpeg/nodes.py @@ -235,7 +235,7 @@ class Node(KwargReprNode): class FilterableStream(Stream): def __init__(self, upstream_node, upstream_label, upstream_selector=None): super(FilterableStream, self).__init__( - upstream_node, upstream_label, {InputNode, FilterNode}, upstream_selector + upstream_node, upstream_label, {HeaderNode, InputNode, FilterNode}, upstream_selector ) @@ -243,11 +243,11 @@ class FilterableStream(Stream): class InputNode(Node): """InputNode type""" - def __init__(self, name, args=[], kwargs={}): + def __init__(self, name, stream: Stream = None, args=[], kwargs={}): super(InputNode, self).__init__( - stream_spec=None, + stream_spec=stream, name=name, - incoming_stream_types={}, + incoming_stream_types={FilterableStream}, outgoing_stream_type=FilterableStream, min_inputs=0, max_inputs=0, @@ -343,6 +343,23 @@ class MergeOutputsNode(Node): # noinspection PyMethodOverriding +class HeaderNode (Node): + + def __init__(self, name: str = None, args=[], kwargs={}): + super(HeaderNode, self).__init__( + stream_spec=None, + name=name, + incoming_stream_types={}, + outgoing_stream_type=FilterableStream, + min_inputs=0, + max_inputs=0, + args=args, + kwargs=kwargs, + ) + +# noinspection PyMethodOverriding + + class GlobalNode(Node): def __init__(self, stream, name, args=[], kwargs={}): super(GlobalNode, self).__init__( diff --git a/ffmpeg/tests/test_ffmpeg.py b/ffmpeg/tests/test_ffmpeg.py index 279a323..aecf069 100644 --- a/ffmpeg/tests/test_ffmpeg.py +++ b/ffmpeg/tests/test_ffmpeg.py @@ -135,6 +135,20 @@ def test_global_args(): ] +def test_header_args(): + out_file = ( + ffmpeg.header(thread_queue_size='512') + .input("input.mp4") + .output("output.mp4") + ) + assert out_file.get_args() == [ + '-thread_queue_size=512', + '-i', + 'input.mp4', + 'output.mp4', + ] + + def _get_simple_example(): return ffmpeg.input(TEST_INPUT_FILE1).output(TEST_OUTPUT_FILE1) @@ -697,7 +711,7 @@ def test_pipe(): out_data = p.stdout.read() assert len(out_data) == frame_size * (frame_count - start_frame) - assert out_data == in_data[start_frame * frame_size :] + assert out_data == in_data[start_frame * frame_size:] def test__probe(): diff --git a/requirements.txt b/requirements.txt index f8b347e..1034d9d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,7 @@ alabaster==0.7.12 atomicwrites==1.3.0 attrs==19.1.0 +autopep8==1.4.4 Babel==2.7.0 certifi==2019.3.9 chardet==3.0.4