Repository: dask/dask

Issue:
Add Base.to_graphviz
Add function to return `graphviz` instance created from dask graph for below reasons:
- When using IPython, `.visualize` outputs unnecessary image file
- Sometimes we want to modify graphviz instance directly



Discussion:
I'm not too keen on having this as a method on collections, seems a little extraneous. `dask.dot.to_graphviz` could be changed to accept either a dask graph or an instance of `Base`. Then this would work:

``` python
>>> from dask.dot import to_graphviz
>>> g = to_graphviz(some_dask_collection)
```

Perhaps temporary files could be used by `visualize` when no path name is specified. This would still allow displaying through IPython notebook, but wouldn't clutter your directory with `mydask.*` files.

Thanks. On current master, specifying `filename=None` results in `TypeError`.

```
d.visualize(filename=None)
TypeError: unsupported operand type(s) for +: 'NoneType' and 'str'
```

I think this can be fixed to:
1. Return image using temporary file
2. Return image itself without saving a file
3. Return graphiz instance

Maybe option 2 is prefereable? (I want 3 in future, but it should be done via different option...)

Historically we've defaulted to a default filename, `'mydask'`.


Resolving patch:
diff --git a/dask/dot.py b/dask/dot.py
--- a/dask/dot.py
+++ b/dask/dot.py
@@ -6,6 +6,7 @@
 from graphviz import Digraph
 
 from .core import istask, get_dependencies, ishashable
+from .compatibility import BytesIO
 
 
 def task_label(task):
@@ -132,19 +133,35 @@ def to_graphviz(dsk, data_attributes=None, function_attributes=None):
 
 def dot_graph(dsk, filename='mydask', **kwargs):
     g = to_graphviz(dsk, **kwargs)
-    g.save(filename + '.dot')
 
-    try:
-        check_call('dot -Tpdf {0}.dot -o {0}.pdf'.format(filename), shell=True)
-        check_call('dot -Tpng {0}.dot -o {0}.png'.format(filename), shell=True)
-    except CalledProcessError:
-        raise RuntimeError(
-            "Please install The `dot` utility from graphviz:\n"
-            "  Debian:  sudo apt-get install graphviz\n"
-            "  Mac OSX: brew install graphviz\n"
-            "  Windows: http://www.graphviz.org/Download..php")  # pragma: no cover
-    try:
-        from IPython.display import Image
-        return Image(filename + '.png')
-    except ImportError:
-        pass
+    if filename is not None:
+        g.save(filename + '.dot')
+
+        try:
+            check_call('dot -Tpdf {0}.dot -o {0}.pdf'.format(filename),
+                       shell=True)
+            check_call('dot -Tpng {0}.dot -o {0}.png'.format(filename),
+                       shell=True)
+
+        except CalledProcessError:
+            msg = ("Please install The `dot` utility from graphviz:\n"
+                   "  Debian:  sudo apt-get install graphviz\n"
+                   "  Mac OSX: brew install graphviz\n"
+                   "  Windows: http://www.graphviz.org/Download..php")
+            raise RuntimeError(msg)  # pragma: no cover
+
+        try:
+            from IPython.display import Image
+            return Image(filename + '.png')
+        except ImportError:
+            pass
+
+    else:
+        try:
+            from IPython.display import Image
+            s = BytesIO()
+            s.write(g.pipe(format='png'))
+            s.seek(0)
+            return Image(s.read())
+        except ImportError:
+            pass


Test patch:
diff --git a/dask/tests/test_dot.py b/dask/tests/test_dot.py
--- a/dask/tests/test_dot.py
+++ b/dask/tests/test_dot.py
@@ -8,6 +8,7 @@
 pytest.importorskip("graphviz")
 
 from dask.dot import dot_graph, task_label, label, to_graphviz
+from IPython.display import Image
 
 # Since graphviz doesn't store a graph, we need to parse the output
 label_re = re.compile('.*\[label=(.*?) shape=.*\]')
@@ -75,9 +76,20 @@ def test_dot_graph():
     fn = 'test_dot_graph'
     fns = [fn + ext for ext in ['.png', '.pdf', '.dot']]
     try:
-        dot_graph(dsk, filename=fn)
+        i = dot_graph(dsk, filename=fn)
         assert all(os.path.exists(f) for f in fns)
+        assert isinstance(i, Image)
     finally:
         for f in fns:
             if os.path.exists(f):
                 os.remove(f)
+
+    fn = 'mydask' # default, remove existing files
+    fns = [fn + ext for ext in ['.png', '.pdf', '.dot']]
+    for f in fns:
+        if os.path.exists(f):
+            os.remove(f)
+    i = dot_graph(dsk, filename=None)
+    assert all(not os.path.exists(f) for f in fns)
+    assert isinstance(i, Image)
+