Ensure that graphviz is an optional dependency

pystorm · Sep 2, 2016 · 5e94744 · 5e94744
1 parent 30c6d70
commit 5e94744
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/streamparse/cli/visualize.py b/streamparse/cli/visualize.py
@@ -15,9 +15,9 @@
 
 try:
     import graphviz
+    HAVE_GRAPHVIZ = True
 except:
-    raise ImportError('The visualize command requires the `graphviz` Python '
-                      'library and `graphviz` system library to be installed.')
+    HAVE_GRAPHVIZ = False
 
 
 IPYTHON_IMAGE_FORMATS = frozenset(['jpeg', 'png'])
@@ -26,6 +26,10 @@
 
 def to_graphviz(topology_class, node_attr=None, edge_attr=None, **kwargs):
     """Convert a Topology into a DiGraph"""
+    if not HAVE_GRAPHVIZ:
+        raise ImportError('The visualize command requires the `graphviz` Python'
+                          ' library and `graphviz` system library to be '
+                          'installed.')
     attributes = {'fontsize': '16',
                   'fontcolor': 'white',
                   'bgcolor': '#333333',