Add zip_with

Author: strokyl

toolz/curried/__init__.py

@@ -51,6 +51,7 @@
     second,
     thread_first,
     thread_last,
+    zip_dict
 )
 from .exceptions import merge, merge_with
 

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,13 @@ 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))

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'