Repository: dask/dask
Pull request #4165 (2018-11-01 09:37:45)

Add traverse support for dataclasses
`Dataclasses` are not traversed right now. As I see them replacing pure `dicts` in quite some cases I think it would be nice to handle them as well.

No tests yet added, wanted to check first if you guys are interested.

- [x] Tests added / passed
- [x] Passes `flake8 dask`


aberres: Hm, seems as if this is incomplete.

While threaded execution works using processes fails when it tries to pickle `Delayed` instances.
Seems as if I am something.
aberres: Needed to scratch that itch... 

Seems to work with the latest update.
mrocklin: @jcrist can I ask you to review and manage this proposal?
aberres: The test in `test_base.py` has a `try: ... except ImportError: ` as otherwise the test fails due to the `dataclass` not coming from a module level scope.

Patch:
diff --git a/dask/base.py b/dask/base.py
index 7cb0c660f..cf2cb3cc7 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -13,7 +13,8 @@ import uuid
 from toolz import merge, groupby, curry, identity
 from toolz.functoolz import Compose
 
-from .compatibility import long, unicode, Iterator
+from .compatibility import (apply, long, unicode, Iterator, is_dataclass,
+                            dataclass_fields)
 from .context import thread_state
 from .core import flatten, quote, get as simple_get
 from .hashing import hash_buffer_hex
@@ -264,6 +265,10 @@ def unpack_collections(*args, **kwargs):
             elif typ is dict:
                 tsk = (dict, [[_unpack(k), _unpack(v)]
                               for k, v in expr.items()])
+            elif is_dataclass(expr):
+                tsk = (apply, typ, (), (dict,
+                       [[f.name, _unpack(getattr(expr, f.name))] for f in
+                        dataclass_fields(expr)]))
             else:
                 return expr
 
diff --git a/dask/compatibility.py b/dask/compatibility.py
index 25adaa1aa..9c9b9c32a 100644
--- a/dask/compatibility.py
+++ b/dask/compatibility.py
@@ -307,7 +307,7 @@ def getargspec(func):
     """Version of inspect.getargspec that works with partial and warps."""
     if isinstance(func, functools.partial):
         return getargspec(func.func)
- 
+
     func = getattr(func, '__wrapped__', func)
     if isinstance(func, type):
         return _getargspec(func.__init__)
@@ -337,3 +337,14 @@ def bind_method(cls, name, func):
         setattr(cls, name, types.MethodType(func, None, cls))
     else:
         setattr(cls, name, func)
+
+
+try:
+    from dataclasses import is_dataclass, fields as dataclass_fields
+
+except ImportError:
+    def is_dataclass(x):
+        return False
+
+    def dataclass_fields(x):
+        return []
diff --git a/dask/delayed.py b/dask/delayed.py
index ff3cbe358..aa2cde5d3 100644
--- a/dask/delayed.py
+++ b/dask/delayed.py
@@ -12,7 +12,8 @@ except ImportError:
 from . import config, threaded
 from .base import is_dask_collection, dont_optimize, DaskMethodsMixin
 from .base import tokenize as _tokenize
-from .compatibility import apply, Iterator
+from .compatibility import apply, Iterator, is_dataclass, dataclass_fields
+
 from .core import quote
 from .context import globalmethod
 from .optimization import cull
@@ -91,6 +92,12 @@ def to_task_dask(expr):
         args, dsk = to_task_dask([[k, v] for k, v in expr.items()])
         return (dict, args), dsk
 
+    if is_dataclass(expr):
+        args, dsk = to_task_dask([[f.name, getattr(expr, f.name)] for f in
+                                  dataclass_fields(expr)])
+
+        return (apply, typ, (), (dict, args)), dsk
+
     if typ is slice:
         args, dsk = to_task_dask([expr.start, expr.stop, expr.step])
         return (slice,) + tuple(args), dsk


Test patch:
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
index ab981e41d..12a867206 100644
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -346,13 +346,21 @@ def test_is_dask_collection():
     assert not is_dask_collection(DummyCollection)
 
 
+try:
+    import dataclasses
+    # Avoid @dataclass decorator as Python < 3.7 fail to interpret the type hints
+    ADataClass = dataclasses.make_dataclass('ADataClass', [('a', int)])
+except ImportError:
+    dataclasses = None
+
+
 def test_unpack_collections():
     a = delayed(1) + 5
     b = a + 1
     c = a + 2
 
     def build(a, b, c, iterator):
-        return (a, b,               # Top-level collections
+        t = (a, b,                  # Top-level collections
                 {'a': a,            # dict
                  a: b,              # collections as keys
                  'b': [1, 2, [b]],  # list
@@ -361,6 +369,11 @@ def test_unpack_collections():
                  'e': {a, 2, 3}},   # set
                 iterator)           # Iterator
 
+        if dataclasses is not None:
+            t[2]['f'] = ADataClass(a=a)
+
+        return t
+
     args = build(a, b, c, (i for i in [a, b, c]))
 
     collections, repack = unpack_collections(*args)
diff --git a/dask/tests/test_delayed.py b/dask/tests/test_delayed.py
index 4a4c888fa..c887f29ca 100644
--- a/dask/tests/test_delayed.py
+++ b/dask/tests/test_delayed.py
@@ -94,6 +94,22 @@ def test_delayed():
     assert a.key in b.dask
 
 
+def test_delayed_with_dataclass():
+    dataclasses = pytest.importorskip("dataclasses")
+
+    # Avoid @dataclass decorator as Python < 3.7 fail to interpret the type hints
+    ADataClass = dataclasses.make_dataclass('ADataClass', [('a', int)])
+
+    literal = dask.delayed(3)
+    with_class = dask.delayed({"a": ADataClass(a=literal)})
+
+    def return_nested(obj):
+        return obj["a"].a
+    final = delayed(return_nested)(with_class)
+
+    assert final.compute() == 3
+
+
 def test_operators():
     a = delayed([1, 2, 3])
     assert a[0].compute() == 1