pmap should support multiple iterables

gbatchy/utils.py

@@ -1,5 +1,6 @@
 from collections import deque, OrderedDict
 from gevent import Timeout, iwait as _gevent_iwait, wait as _gevent_wait, pool as _gevent_pool, queue as _gevent_queue, sleep
+from itertools import izip_longest
 try:
     from peak.util.proxies import LazyProxy
 except ImportError:
@@ -32,8 +33,8 @@ def pget(lst):
     return [x.get() for x in lst]
 
 @batch_context
-def pmap(fn, items, **kwargs):
-    return pget(spawn(fn, i, **kwargs) for i in items)
+def pmap(fn, *seqs, **kwargs):
+    return pget(spawn(fn, *args, **kwargs) for args in izip_longest(*seqs))
 
 @batch_context
 def pmap_unordered(fn, items, **kwargs):
@@ -433,4 +434,4 @@ def pmap(self, *args, **kwargs):
         return list(self.imap(*args, **kwargs))
     map = pmap
 
-    
+

tests/batch_tests.py

@@ -187,6 +187,10 @@ def test():
         test()  # shouldn't hang.
 
     def test_utils(self):
+        @batched()
+        def add(x, y):
+            return x + y
+
         @batched()
         def add_n(args_list):
             return [args[0] + kwargs.get('n', 1)
@@ -202,6 +206,8 @@ def test():
             self.assertEquals([3,4,5], pmap(add_n, [1,2,3], n=2))
             self.assertEquals([2], pfilter(only_even, [1,2,3]))
 
+            self.assertEquals([11,13,23], pmap([7,11,10], [4,2,13]))
+
             self.assertEquals([2,3,4], sorted(pmap_unordered(add_n, [1,2,3])))
             self.assertEquals([3,4,5], sorted(pmap_unordered(add_n, [1,2,3], n=2)))
             self.assertEquals([2], list(pfilter_unordered(only_even, [1,2,3])))