Overload HLG values method (dask#4918)

jrbourbeau · jcrist · commit 46aef5870482 · 2019-06-13T11:04:47.000-05:00
* Overload HLG values method

* Return lists for keys, values, and items

* Add tests for keys and items
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
@@ -143,16 +143,24 @@ def __len__(self):
         return sum(1 for _ in self)
 
     def items(self):
+        items = []
         seen = set()
         for d in self.layers.values():
             for key in d:
                 if key not in seen:
                     seen.add(key)
-                    yield (key, d[key])
+                    items.append((key, d[key]))
+        return items
 
     def __iter__(self):
         return toolz.unique(toolz.concat(self.layers.values()))
 
+    def keys(self):
+        return [key for key, _ in self.items()]
+
+    def values(self):
+        return [value for _, value in self.items()]
+
     @classmethod
     def merge(cls, *graphs):
         layers = {}
diff --git a/dask/tests/test_highgraph.py b/dask/tests/test_highgraph.py
@@ -26,3 +26,17 @@ def test_basic():
     hg = HighLevelGraph(layers, dependencies)
 
     assert dict(hg) == {'x': 1, 'y': (inc, 'x')}
+
+
+def test_keys_values_items_methods():
+    a = da.ones(10, chunks=(5,))
+    b = a + 1
+    c = a + 2
+    d = b + c
+    hg = d.dask
+
+    keys, values, items = hg.keys(), hg.values(), hg.items()
+    assert all(isinstance(i, list) for i in [keys, values, items])
+    assert keys == [i for i in hg]
+    assert values == [hg[i] for i in hg]
+    assert items == [(k, v) for k, v in zip(keys, values)]
