diff --git a/ffmpeg/__init__.py b/ffmpeg/__init__.py index 1953609..9f1bed7 100644 --- a/ffmpeg/__init__.py +++ b/ffmpeg/__init__.py @@ -3,4 +3,5 @@ from . import _filters, _ffmpeg, _run from ._filters import * from ._ffmpeg import * from ._run import * -__all__ = _filters.__all__ + _ffmpeg.__all__ + _run.__all__ +from ._view import * +__all__ = _filters.__all__ + _ffmpeg.__all__ + _run.__all__ + _view.__all__ diff --git a/ffmpeg/graph.py b/ffmpeg/_view.py similarity index 80% rename from ffmpeg/graph.py rename to ffmpeg/_view.py index 9e220df..aec939c 100644 --- a/ffmpeg/graph.py +++ b/ffmpeg/_view.py @@ -2,8 +2,8 @@ from __future__ import unicode_literals from .dag import get_outgoing_edges from ._run import topo_sort -import graphviz import os +import tempfile from ffmpeg.nodes import ( InputNode, @@ -15,10 +15,20 @@ from ffmpeg.nodes import ( @operator() def view(*downstream_nodes, **kwargs): + try: + import graphviz + except ImportError: + raise ImportError('failed to import graphviz; please make sure graphviz is installed (e.g. `pip install ' + 'graphviz`)') + + filename = kwargs.pop('filename', None) + show_labels = kwargs.pop('show_labels', False) + if filename is None: + filename = tempfile.mktemp() + sorted_nodes, outgoing_edge_maps = topo_sort(downstream_nodes) graph = graphviz.Digraph() graph.attr(rankdir='LR') - show_labels = kwargs.pop('show_labels', False) if len(kwargs.keys()) != 0: raise ValueError('Invalid kwargs key(s): {}'.format(', '.join(kwargs.keys()))) @@ -44,4 +54,10 @@ def view(*downstream_nodes, **kwargs): downstream_node_id = str(hash(edge.downstream_node)) graph.edge(upstream_node_id, downstream_node_id, **kwargs) - graph.view('graph.png') + graph.view(filename) + + + +__all__ = [ + 'view', +]