FIX(PLOT.TC): TC was always testing PNG, ...

+ retorfitted to try all available formats.
+ list of forbidden formats based on my failres

Author: ankostis

graphkit/base.py

@@ -176,8 +176,8 @@ def plot(self, filename=None, show=False,
         """
         :param str filename:
             Write diagram into a file.
-            The extension must be one of: ``.png .dot .jpg .jpeg .pdf .svg``
-            Prefer ``.pdf`` or ``.svg`` to see solution-values in tooltips.
+            Common extensions are ``.png .dot .jpg .jpeg .pdf .svg``
+            call :func:`network.supported_plot_formats()` for more.
         :param boolean show:
             If it evaluates to true, opens the  diagram in a  matplotlib window.
         :param inputs:

graphkit/network.py

@@ -383,8 +383,8 @@ def plot(self, filename=None, show=False,
 
         :param str filename:
             Write diagram into a file.
-            The extension must be one of: ``.png .dot .jpg .jpeg .pdf .svg``
-            Prefer ``.pdf`` or ``.svg`` to see solution-values in tooltips.
+            Common extensions are ``.png .dot .jpg .jpeg .pdf .svg``
+            call :func:`network.supported_plot_formats()` for more.
         :param boolean show:
             If it evaluates to true, opens the  diagram in a  matplotlib window.
         :param inputs:
@@ -454,15 +454,10 @@ def get_data_node(name, graph):
     return None
 
 
-def supported_plot_writers():
-    return {
-        ".png": lambda gplot: gplot.create_png(),
-        ".dot": lambda gplot: gplot.to_string(),
-        ".jpg": lambda gplot: gplot.create_jpeg(),
-        ".jpeg": lambda gplot: gplot.create_jpeg(),
-        ".pdf": lambda gplot: gplot.create_pdf(),
-        ".svg": lambda gplot: gplot.create_svg(),
-    }
+def supported_plot_formats():
+    import pydot
+
+    return [".%s" % f for f in pydot.Dot().formats]
 
 
 def plot_graph(graph, filename=None, show=False, steps=None,
@@ -496,8 +491,8 @@ def plot_graph(graph, filename=None, show=False, steps=None,
         what to plot
     :param str filename:
         Write diagram into a file.
-        The extension must be one of: ``.png .dot .jpg .jpeg .pdf .svg``
-        Prefer ``.pdf`` or ``.svg`` to see solution-values in tooltips.
+        Common extensions are ``.png .dot .jpg .jpeg .pdf .svg``
+        call :func:`network.supported_plot_formats()` for more.
     :param boolean show:
         If it evaluates to true, opens the  diagram in a  matplotlib window.
     :param steps:
@@ -603,16 +598,15 @@ def get_node_name(a):
 
     # save plot
     if filename:
+        formats = supported_plot_formats()
         _basename, ext = os.path.splitext(filename)
-        writers = supported_plot_writers()
-        plot_writer = supported_plot_writers().get(ext.lower())
-        if not plot_writer:
+        if not ext.lower() in formats:
             raise ValueError(
                 "Unknown file format for saving graph: %s"
                 "  File extensions must be one of: %s"
-                % (ext, ' '.join(writers)))
-        with open(filename, "wb") as fh:
-            fh.write(plot_writer(g))
+                % (ext, " ".join(formats)))
+
+        g.write(filename, format=ext.lower()[1:])
 
     # display graph via matplotlib
     if show:

test/test_graphkit.py

@@ -116,27 +116,41 @@ def test_plotting():
     inputs = {'a': 1, 'b1': 2}
     solution=pipeline(inputs)
 
-    for ext in network.supported_plot_writers():
-        tdir = tempfile.mkdtemp(suffix=ext)
-        png_file = osp.join(tdir, "workflow.png")
-        pipeline.plot(png_file, inputs=inputs, solution=solution, outputs=['asked', 'b1'])
-        try:
-            assert osp.exists(png_file)
-        finally:
-            shutil.rmtree(tdir, ignore_errors=True)
+    # ...not working on my PC ...
+    forbidden_formats = ".dia .hpgl .mif .pcl .pic .vtx .xlib".split()
+    tdir = tempfile.mkdtemp()
+    counter = 0
+    try:
+        for ext in network.supported_plot_formats():
+            if ext in forbidden_formats:
+                continue
+
+            counter += 1
+            fpath = osp.join(tdir, "workflow-%i%s" % (counter, ext))
+            pipeline.plot(fpath, inputs=inputs, solution=solution, outputs=['asked', 'b1'])
+            assert osp.exists(fpath)
+
+            counter += 1
+            fpath = osp.join(tdir, "workflow-%i%s" % (counter, ext))
+            pipeline.plot(fpath)
+            assert osp.exists(fpath)
+    finally:
+        shutil.rmtree(tdir, ignore_errors=True)
+
     try:
         pipeline.plot('bad.format')
         assert False, "Should had failed writting arbitrary file format!"
     except ValueError as ex:
         assert "Unknown file format" in str(ex)
 
         ## Check help msg lists all siupported formats
-        for ext in network.supported_plot_writers():
+        for ext in network.supported_plot_formats():
             assert ext in str(ex)
 
 
 def test_plotting_docstring():
-    for ext in network.supported_plot_writers():
+    common_formats = ".png .dot .jpg .jpeg .pdf .svg".split()
+    for ext in common_formats:
         assert ext in network.plot_graph.__doc__