blackteay/ffmpeg-python
1
0
Fork 0
mirror of https://github.com/kkroening/ffmpeg-python.git synced 2025-04-06 04:15:44 +08:00
Code Issues Projects Releases Wiki Activity

#18: fix to use latest #17 changes

Browse Source
This commit is contained in:
Karl Kroening 2017-07-06 02:24:40 -06:00
parent 656d9fa7f7
commit 1955547202
1 changed files with 32 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

48
ffmpeg/_view.py
View File

@ -6,15 +6,28 @@ import os
import tempfile import tempfile
from ffmpeg.nodes import ( from ffmpeg.nodes import (
FilterNode,
InputNode, InputNode,
OutputNode, OutputNode,
FilterNode, Stream,
operator, stream_operator,
) )
@operator() def _get_node_color(node):
def view(*downstream_nodes, **kwargs): if isinstance(node, InputNode):
color = '#99cc00'
elif isinstance(node, OutputNode):
color = '#99ccff'
elif isinstance(node, FilterNode):
color = '#ffcc00'
else:
color = None
return color
@stream_operator()
def view(*streams, **kwargs):
try: try:
import graphviz import graphviz
except ImportError: except ImportError:
@ -22,11 +35,17 @@ def view(*downstream_nodes, **kwargs):
'graphviz`)') 'graphviz`)')
filename = kwargs.pop('filename', None) filename = kwargs.pop('filename', None)
show_labels = kwargs.pop('show_labels', False) show_labels = kwargs.pop('show_labels', True)
if filename is None: if filename is None:
filename = tempfile.mktemp() filename = tempfile.mktemp()
sorted_nodes, outgoing_edge_maps = topo_sort(downstream_nodes) nodes = []
for stream in streams:
if not isinstance(stream, Stream):
raise TypeError('Expected Stream; got {}'.format(type(stream)))
nodes.append(stream.node)
sorted_nodes, outgoing_edge_maps = topo_sort(nodes)
graph = graphviz.Digraph() graph = graphviz.Digraph()
graph.attr(rankdir='LR') graph.attr(rankdir='LR')
if len(kwargs.keys()) != 0: if len(kwargs.keys()) != 0:
@ -36,20 +55,17 @@ def view(*downstream_nodes, **kwargs):
name = node.name name = node.name
if '_kwargs' in dir(node) and 'filename' in node._kwargs: if '_kwargs' in dir(node) and 'filename' in node._kwargs:
name = os.path.basename(node._kwargs['filename']) name = os.path.basename(node._kwargs['filename'])
if isinstance(node, InputNode): color = _get_node_color(node)
color = '#99cc00'
elif isinstance(node, OutputNode):
color = '#99ccff'
elif isinstance(node, FilterNode):
color = '#ffcc00'
else:
color = None
graph.node(str(hash(node)), name, shape='box', style='filled', fillcolor=color) graph.node(str(hash(node)), name, shape='box', style='filled', fillcolor=color)
outgoing_edge_map = outgoing_edge_maps.get(node, {}) outgoing_edge_map = outgoing_edge_maps.get(node, {})
for edge in get_outgoing_edges(node, outgoing_edge_map): for edge in get_outgoing_edges(node, outgoing_edge_map):
kwargs = {} kwargs = {}
if show_labels: if show_labels and (edge.upstream_label is not None or edge.downstream_label is not None):
kwargs['label'] = '{} -> {}'.format(edge.upstream_label, edge.downstream_label) upstream_label = edge.upstream_label if edge.upstream_label is not None else ''
downstream_label = edge.downstream_label if edge.downstream_label is not None else ''
kwargs['label'] = '{} {}'.format(upstream_label, downstream_label)
upstream_node_id = str(hash(edge.upstream_node)) upstream_node_id = str(hash(edge.upstream_node))
downstream_node_id = str(hash(edge.downstream_node)) downstream_node_id = str(hash(edge.downstream_node))
graph.edge(upstream_node_id, downstream_node_id, **kwargs) graph.edge(upstream_node_id, downstream_node_id, **kwargs)