Repository: dask/dask

Issue:
dask.imperative.do kwargs support
Hey Blaze team! I've just started using Blaze/Dask/Odo, and I love your work.

I came across this issue recently--I would expect the following behavior from dask:

``` python
>>> from dask import do
>>> @do
>>> def mysum(*args, **kwargs):
>>>     return sum(args) + sum(kwargs.values())
>>> ten = mysum(1,2,3, four=mysum(2,2))
>>> ten.compute() # currently returns a Value object
10
```

Currently, however, it returns a value object, from which it looks like the kwargs in mysum aren't being considered as dask-able. I do get the expected result when I monkey-patch dask.imperative.applyfunc with the following:

``` python
def applyfunc(func, args, kwargs, pure=False):
    """Create a Value by applying a function to args.
    Given a function and arguments, return a Value that represents the result
    of that computation."""

    args, dasks = unzip(map(to_task_dasks, args), 2)
    if kwargs:
        # change tokenize implementation so we don't need this?
        assert 'pure' not in kwargs, 'reserved keyword' 
        task_kwvals, kwdasks = unzip(map(to_task_dasks, kwargs.values()), 2)
        kwargs = dict(zip(kwargs.keys(), task_kwvals))
        dasks += kwdasks
    dasks = flat_unique(dasks)
    name = tokenize(func, *args, pure=pure, **kwargs)
    if kwargs:
        # I have no idea why this works, it could probably be simpler
        dasks.append({name: (apply, func, args,
                             (dict, (list, list(map(list, kwargs.items())))))})
    else:
        dasks.append({name: (func,) + args})
    return Value(name, dasks)
```

I tried it with some additional cleverness to keep the functools.partial speedup when possible, but I think it may overcomplicate things:

``` python
    if kwargs:
        split_kwargs = groupby(lambda item: isinstance(item[1], Value),
                               kwargs.items())
        func = partial(func, **dict(split_kwargs.get(False, [])))
        task_kwargs = dict(split_kwargs.get(True, []))
        # change tokenize implementation so we don't need this?
        assert 'pure' not in task_kwargs, 'reserved keyword' 
        task_kwvals, kwdasks = unzip(map(to_task_dasks, task_kwargs.values()), 2)
        kwargs = dict(zip(task_kwargs.keys(), task_kwvals))
        dasks += kwdasks
```

...but the weird reservation for the 'pure' keyword still bugs me a bit. Any ideas for what the least invasive way would be to fix that would be? Or is this entire idea somehow antithetical?

Thanks!

PS: Apologies that this isn't in a pull-request



Discussion:

Resolving patch:
diff --git a/dask/imperative.py b/dask/imperative.py
--- a/dask/imperative.py
+++ b/dask/imperative.py
@@ -1,5 +1,5 @@
 import operator
-from functools import partial, wraps
+from functools import wraps
 from itertools import chain, count
 from collections import Iterator
 
@@ -114,11 +114,15 @@ def applyfunc(func, args, kwargs, pure=False):
     of that computation."""
 
     args, dasks = unzip(map(to_task_dasks, args), 2)
-    dasks = flat_unique(dasks)
     if kwargs:
-        func = partial(func, **kwargs)
-    name = tokenize(func, *args, pure=pure)
-    dasks.append({name: (func,) + args})
+        dask_kwargs, dasks2 = to_task_dasks(kwargs)
+        dasks = dasks + (dasks2,)
+        task = (apply, func, (list, list(args)), dask_kwargs)
+    else:
+        task = (func,) + args
+    name = tokenize(*task, pure=pure)
+    dasks = flat_unique(dasks)
+    dasks.append({name: task})
     return Value(name, dasks)
 
 
@@ -240,7 +244,7 @@ def __hash__(self):
         return hash(self.key)
 
     def __dir__(self):
-        return list(self.__dict__.keys())
+        return dir(type(self))
 
     def __getattr__(self, attr):
         if not attr.startswith('_'):


Test patch:
diff --git a/dask/tests/test_imperative.py b/dask/tests/test_imperative.py
--- a/dask/tests/test_imperative.py
+++ b/dask/tests/test_imperative.py
@@ -164,6 +164,17 @@ def test_pure():
     assert myrand().key != myrand().key
 
 
+def test_kwargs():
+    def mysum(a, b, c=(), **kwargs):
+        return a + b + sum(c) + sum(kwargs.values())
+    dmysum = do(mysum)
+    ten = dmysum(1, 2, c=[value(3), 0], four=dmysum(2,2))
+    assert ten.compute() == 10
+    dmysum = do(mysum, pure=True)
+    ten = dmysum(1, 2, c=[value(3), 0], four=dmysum(2,2))
+    assert ten.compute() == 10
+
+
 da = pytest.importorskip('dask.array')
 import numpy as np