Repository: dask/dask

Issue:
`dask.compute` on nested Python objects
I really like using `dask.compute` to calculate many things in a single pass over the data.  I sometimes even do so programmatically--e.g., from using results from a previous pass over the data.  The single list format via `*args` in `compute` is somewhat limiting.  I want a `compute` that also computes elements from nested dicts and lists.  I hacked/whipped up the following to do this that uses `dask.compute`:
```python
# Extend `dask.compute` to work on nested data structures.
from toolz import concat
from dask import compute

def _flatten_compute(vals):
    if isinstance(vals, list):
        inner = list(concat(map(_flatten_compute, vals)))
        return ['list', len(inner)] +  inner
    elif isinstance(vals, dict):
        inner = []
        for k, v in vals.items():
            inner.append(k)
            inner.extend(_flatten_compute(v))
        return ['dict', len(inner)] +  inner
    else:
        return ['item', vals]


def _rebuild_compute(results, start=0):
    if results[start] == 'list':
        rv = []
        stop = start + 2 + results[start + 1]
        i = start + 2
        while i < stop:
            val = results[i]
            if val == 'item':
                rv.append(results[i + 1])
                i += 2
            else:  # list or dict
                rv.append(_rebuild_compute(results, start=i))
                i += 2 + results[i + 1]
        return rv
    elif results[start] == 'dict':
        rv = {}
        stop = start + 2 + results[start + 1]
        i = start + 2
        while i < stop:
            key = results[i]
            val = results[i + 1]
            if val == 'item':
                rv[key] = results[i + 2]
                i += 3
            else:  # list or dict
                rv[key] = _rebuild_compute(results, start=i + 1)
                i += 3 + results[i + 2]
        return rv
    else:
        raise ValueError('Does not compute: %s %s' % (start, results))


def compute_nested(*args, **kwargs):
    """Extended ``dask.compute`` to work on nested data structures.
    
    This will compute (potentially nested) elements in lists and
    values in dicts.  For example:
    
    >>> compute_nested({'total': bag.count(), 'truthy': bag.filter(None).count()})
    ({'total': 100, 'truthy': 75},)
    """    
    flattened_args = _flatten_compute(list(args))
    results = compute(*flattened_args, **kwargs)
    return _rebuild_compute(results)
```
Thoughts?  Should I take the effort to make a PR to add this to `dask`?


Discussion:
This already exists as `dask.delayed.compute`. I'm not sure if I think this should replace `dask.compute` though - there's some benefit to cheaply passing through non-dask objects (which is no longer cheap in `dask.delayed.compute`).
Oh, I had no idea that existed!  Thanks.  It's also a much cleaner implementation.  `dask.delayed` ftw!
It appears that it's only referenced in the `dask.delayed` api docs (http://dask.pydata.org/en/latest/delayed-api.html#dask.delayed.compute), so it may be non-obvious. Thoughts (or PR:) ) on how to improve the docs (if needed)? 
Personally, I wouldn't mind of `dask.compute` was replaced by `dask.delayed.compute`.
I'm not 100% against it. My one issues is that currently `compute` passes through non-dask objects very cheaply (just an `isinstance` check). Switching to `dask.delayed.compute` results in a more expensive crawl of the objects beforehand if the object is one of the builtin collections (tuple, dict, set, list). This may not be a big deal though - in practice I only know of `dask-glm` that might call `compute` on non-dask objects.

I also see that `dask.delayed.compute` converts all arguments to `delayed` - really it should only do that for things that are not subclasses of `Base`. This should be fixed regardless of the decision above.
Rather than having two similar functions (and one very hidden to the average user), I think it's preferable to have a keyword argument control behavior, such as `recurse=True`.

Resolving patch:
diff --git a/dask/base.py b/dask/base.py
--- a/dask/base.py
+++ b/dask/base.py
@@ -1,6 +1,6 @@
 from __future__ import absolute_import, division, print_function
 
-from collections import OrderedDict
+from collections import OrderedDict, Iterator
 import copy
 from functools import partial
 from hashlib import md5
@@ -90,7 +90,7 @@ def compute(self, **kwargs):
         kwargs
             Extra keywords to forward to the scheduler ``get`` function.
         """
-        (result,) = compute(self, **kwargs)
+        (result,) = compute(self, traverse=False, **kwargs)
         return result
 
     @classmethod
@@ -140,9 +140,15 @@ def compute(*args, **kwargs):
     Parameters
     ----------
     args : object
-        Any number of objects. If the object is a dask collection, it's
-        computed and the result is returned. Otherwise it's passed through
-        unchanged.
+        Any number of objects. If it is a dask object, it's computed and the
+        result is returned. By default, python builtin collections are also
+        traversed to look for dask objects (for more information see the
+        ``traverse`` keyword). Non-dask arguments are passed through unchanged.
+    traverse : bool, optional
+        By default dask traverses builtin python collections looking for dask
+        objects passed to ``compute``. For large collections this can be
+        expensive. If none of the arguments contain any dask objects, set
+        ``traverse=False`` to avoid doing this traversal.
     get : callable, optional
         A scheduler ``get`` function to use. If not provided, the default is
         to check the global settings first, and then fall back to defaults for
@@ -161,7 +167,19 @@ def compute(*args, **kwargs):
     >>> b = da.arange(10, chunks=2).mean()
     >>> compute(a, b)
     (45, 4.5)
+
+    By default, dask objects inside python collections will also be computed:
+
+    >>> compute({'a': a, 'b': b, 'c': 1})  # doctest: +SKIP
+    ({'a': 45, 'b': 4.5, 'c': 1},)
     """
+    from dask.delayed import delayed
+    traverse = kwargs.pop('traverse', True)
+    if traverse:
+        args = tuple(delayed(a)
+                     if isinstance(a, (list, set, tuple, dict, Iterator))
+                     else a for a in args)
+
     optimize_graph = kwargs.pop('optimize_graph', True)
     variables = [a for a in args if isinstance(a, Base)]
     if not variables:
diff --git a/dask/delayed.py b/dask/delayed.py
--- a/dask/delayed.py
+++ b/dask/delayed.py
@@ -7,13 +7,13 @@
 
 from toolz import merge, unique, curry, first
 
+from . import base, threaded
+from .base import compute  # noqa: compute import for backwards compatability
+from .compatibility import apply
 from .core import quote
 from .utils import concrete, funcname, methodcaller
-from . import base
-from .compatibility import apply
-from . import threaded
 
-__all__ = ['compute', 'do', 'Delayed', 'delayed']
+__all__ = ['do', 'Delayed', 'delayed']
 
 
 def flat_unique(ls):
@@ -282,27 +282,6 @@ def delayed(obj, name=None, pure=False, nout=None, traverse=True):
 do = delayed
 
 
-def compute(*args, **kwargs):
-    """Evaluate more than one ``Delayed`` at once.
-
-    Note that the only difference between this function and
-    ``dask.base.compute`` is that this implicitly wraps python objects in
-    ``Delayed``, allowing for collections of dask objects to be computed.
-
-    Examples
-    --------
-    >>> a = delayed(1)
-    >>> b = a + 2
-    >>> c = a + 3
-    >>> compute(b, c)  # Compute both simultaneously
-    (3, 4)
-    >>> compute(a, [b, c])  # Works for lists of Delayed
-    (1, [3, 4])
-    """
-    args = [delayed(a) for a in args]
-    return base.compute(*args, **kwargs)
-
-
 def right(method):
     """Wrapper to create 'right' version of operator given left version"""
     def _inner(self, other):


Test patch:
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -330,6 +330,19 @@ def test_compute_with_literal():
     assert compute(5) == (5,)
 
 
+def test_compute_nested():
+    a = delayed(1) + 5
+    b = a + 1
+    c = a + 2
+    assert (compute({'a': a, 'b': [1, 2, b]}, (c, 2)) ==
+            ({'a': 6, 'b': [1, 2, 7]}, (8, 2)))
+
+    res = compute([a, b], c, traverse=False)
+    assert res[0][0] is a
+    assert res[0][1] is b
+    assert res[1] == 8
+
+
 @pytest.mark.skipif('not da')
 @pytest.mark.skipif(sys.flags.optimize == 2,
                     reason="graphviz exception with Python -OO flag")
diff --git a/dask/tests/test_delayed.py b/dask/tests/test_delayed.py
--- a/dask/tests/test_delayed.py
+++ b/dask/tests/test_delayed.py
@@ -103,15 +103,6 @@ def test_delayed_errors():
     pytest.raises(TypeError, lambda: bool(a))
 
 
-def test_compute():
-    a = delayed(1) + 5
-    b = a + 1
-    c = a + 2
-    assert compute(b, c) == (7, 8)
-    assert compute(b) == (7,)
-    assert compute([a, b], c) == ([6, 7], 8)
-
-
 def test_common_subexpressions():
     a = delayed([1, 2, 3])
     res = a[0] + a[0]