Skip to content

Commit

Permalink
Overload HLG values method (#4918)
Browse files Browse the repository at this point in the history
* Overload HLG values method

* Return lists for keys, values, and items

* Add tests for keys and items
  • Loading branch information
jrbourbeau authored and jcrist committed Jun 13, 2019
1 parent f9cd802 commit 46aef58
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 1 deletion.
10 changes: 9 additions & 1 deletion dask/highlevelgraph.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {}
Expand Down
14 changes: 14 additions & 0 deletions dask/tests/test_highgraph.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)]

0 comments on commit 46aef58

Please sign in to comment.