from __future__ import unicode_literals from builtins import str from .dag import get_outgoing_edges from ._run import topo_sort import tempfile from ffmpeg.nodes import ( FilterNode, get_stream_spec_nodes, InputNode, OutputNode, stream_operator, ) _RIGHT_ARROW = '\u2192' def _get_node_color(node): 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(stream_spec, detail=False, filename=None, pipe=False, **kwargs): try: import graphviz except ImportError: raise ImportError( 'failed to import graphviz; please make sure graphviz is installed (e.g. `pip install ' 'graphviz`)' ) show_labels = kwargs.pop('show_labels', True) if pipe and filename is not None: raise ValueError('Can\'t specify both `filename` and `pipe`') elif not pipe and filename is None: filename = tempfile.mktemp() nodes = get_stream_spec_nodes(stream_spec) sorted_nodes, outgoing_edge_maps = topo_sort(nodes) graph = graphviz.Digraph(format='png') graph.attr(rankdir='LR') if len(list(kwargs.keys())) != 0: raise ValueError( 'Invalid kwargs key(s): {}'.format(', '.join(list(kwargs.keys()))) ) for node in sorted_nodes: color = _get_node_color(node) if detail: lines = [node.short_repr] lines += ['{!r}'.format(arg) for arg in node.args] lines += [ '{}={!r}'.format(key, node.kwargs[key]) for key in sorted(node.kwargs) ] node_text = '\n'.join(lines) else: node_text = node.short_repr graph.node( str(hash(node)), node_text, shape='box', style='filled', fillcolor=color ) outgoing_edge_map = outgoing_edge_maps.get(node, {}) for edge in get_outgoing_edges(node, outgoing_edge_map): kwargs = {} up_label = edge.upstream_label down_label = edge.downstream_label up_selector = edge.upstream_selector if show_labels and ( up_label is not None or down_label is not None or up_selector is not None ): if up_label is None: up_label = '' if up_selector is not None: up_label += ":" + up_selector if down_label is None: down_label = '' if up_label != '' and down_label != '': middle = ' {} '.format(_RIGHT_ARROW) else: middle = '' kwargs['label'] = '{} {} {}'.format(up_label, middle, down_label) upstream_node_id = str(hash(edge.upstream_node)) downstream_node_id = str(hash(edge.downstream_node)) graph.edge(upstream_node_id, downstream_node_id, **kwargs) if pipe: return graph.pipe() else: graph.view(filename, cleanup=True) return stream_spec __all__ = ['view']