support dataclasses

mrocklin · Nov 17, 2018 · 48e0031 · 48e0031
1 parent 679d439
commit 48e0031
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/dask/delayed.py b/dask/delayed.py
@@ -106,6 +106,12 @@ def unpack_collections(expr):
         args, collections = unpack_collections([expr.start, expr.stop, expr.step])
         return (slice,) + tuple(args), collections
 
+    if is_dataclass(expr):
+        args, collections = unpack_collections([[f.name, getattr(expr, f.name)] for f in
+                                               dataclass_fields(expr)])
+
+        return (apply, typ, (), (dict, args)), collections
+
     return expr, ()
 
 

diff --git a/dask/tests/test_delayed.py b/dask/tests/test_delayed.py
@@ -107,6 +107,7 @@ def test_delayed_with_dataclass():
 
     def return_nested(obj):
         return obj["a"].a
+
     final = delayed(return_nested)(with_class)
 
     assert final.compute() == 3