diff --git a/toolz/curried/__init__.py b/toolz/curried/__init__.py index 43aeffd4..06c1c368 100644 --- a/toolz/curried/__init__.py +++ b/toolz/curried/__init__.py @@ -51,6 +51,7 @@ second, thread_first, thread_last, + zip_dict ) from .exceptions import merge, merge_with diff --git a/toolz/itertoolz.py b/toolz/itertoolz.py index 941b9a9f..9f109659 100644 --- a/toolz/itertoolz.py +++ b/toolz/itertoolz.py @@ -14,7 +14,7 @@ 'first', 'second', 'nth', 'last', 'get', 'concat', 'concatv', 'mapcat', 'cons', 'interpose', 'frequencies', 'reduceby', 'iterate', 'sliding_window', 'partition', 'partition_all', 'count', 'pluck', - 'join', 'tail', 'diff', 'topk', 'peek', 'random_sample') + 'join', 'tail', 'diff', 'topk', 'peek', 'random_sample', 'zip_dict') def remove(predicate, seq): @@ -1004,3 +1004,14 @@ def random_sample(prob, seq, random_state=None): if not hasattr(random_state, 'random'): random_state = Random(random_state) return filter(lambda _: random_state.random() < prob, seq) + + +def zip_dict(*dicts): + if dicts: + def is_in_all_others_dict(k): + return all(map(lambda d: k in d, dicts[1:])) + + common_keys = filter(is_in_all_others_dict, dicts[0].keys()) + + for k in common_keys: + yield k, tuple(map(lambda d: d[k], dicts)) diff --git a/toolz/tests/test_itertoolz.py b/toolz/tests/test_itertoolz.py index 8a5c2407..412c957f 100644 --- a/toolz/tests/test_itertoolz.py +++ b/toolz/tests/test_itertoolz.py @@ -13,7 +13,7 @@ reduceby, iterate, accumulate, sliding_window, count, partition, partition_all, take_nth, pluck, join, - diff, topk, peek, random_sample) + diff, topk, peek, random_sample, zip_dict) from toolz.compatibility import range, filter from operator import add, mul @@ -123,6 +123,21 @@ def test_isdistinct(): assert isdistinct(iter([1, 2, 1])) is False +def test_zip_dict (): + assert tuple(zip_dict()) == () + assert tuple(zip_dict({0: 0})) == ((0, (0 ,)),) + assert tuple(zip_dict({0: 0}, {1: 1})) == () + assert tuple(zip_dict({0: 0}, {0: 1})) == ((0, (0 , 1 )),) + assert tuple(zip_dict({0: 0}, {0: 0, 1: 1})) == ((0, (0 , 0)),) + assert tuple(zip_dict({0: 1, 1: 2}, {0: 0, 1: 1})) == ((0, (1 , 0)), + (1, (2, 1))) + + assert tuple(zip_dict({-1: 0, 0: 0, 1: 1, 2: 2}, + {-2: 0, 0: 1, 1: 2, 2: 3}, + {-3: 0, 0: 2, 1: 3, 2: 4})) == ((0, (0 , 1, 2)), + (1, (1, 2, 3 )), + (2, (2, 3, 4))) + def test_nth(): assert nth(2, 'ABCDE') == 'C' assert nth(2, iter('ABCDE')) == 'C'