From a3ee8de989ec5949a78aa09c52ed7e2d2cb0c166 2011-12-02 21:55:20
From: MinRK <benjaminrk@gmail.com>
Date: 2011-12-02 21:55:20
Subject: [PATCH] support iterators in view.map

Some objects are not sliceable (e.g. xrange). For these, fall back
on itertools.islice.

tests included

reported on IRC by @juliantaylor

---

diff --git a/IPython/parallel/client/map.py b/IPython/parallel/client/map.py
index 1829d8c..dab382a 100644
--- a/IPython/parallel/client/map.py
+++ b/IPython/parallel/client/map.py
@@ -27,6 +27,7 @@ Authors:
 from __future__ import division
 
 import types
+from itertools import islice
 
 from IPython.utils.data import flatten as utils_flatten
 
@@ -77,9 +78,14 @@ class Map:
             else:
                 lo.append(n*basesize + remainder)
                 hi.append(lo[-1] + basesize)
-
         
-        result = seq[lo[p]:hi[p]]
+        try:
+            result = seq[lo[p]:hi[p]]
+        except TypeError:
+            # some objects (iterators) can't be sliced,
+            # use islice:
+            result = list(islice(seq, lo[p], hi[p]))
+            
         return result
            
     def joinPartitions(self, listOfPartitions):
diff --git a/IPython/parallel/tests/test_lbview.py b/IPython/parallel/tests/test_lbview.py
index 39ac107..00756ef 100644
--- a/IPython/parallel/tests/test_lbview.py
+++ b/IPython/parallel/tests/test_lbview.py
@@ -95,6 +95,17 @@ class TestLoadBalancedView(ClusterTestCase):
         # Ensure that results came in order
         self.assertEquals(astheycame, reference)
         self.assertEquals(amr.result, reference)
+
+    def test_map_iterable(self):
+        """test map on iterables (balanced)"""
+        view = self.view
+        # 101 is prime, so it won't be evenly distributed
+        arr = range(101)
+        # so that it will be an iterator, even in Python 3
+        it = iter(arr)
+        r = view.map_sync(lambda x:x, arr)
+        self.assertEquals(r, list(arr))
+
     
     def test_abort(self):
         view = self.view
diff --git a/IPython/parallel/tests/test_view.py b/IPython/parallel/tests/test_view.py
index d513ca2..81022d8 100644
--- a/IPython/parallel/tests/test_view.py
+++ b/IPython/parallel/tests/test_view.py
@@ -240,6 +240,16 @@ class TestView(ClusterTestCase):
         r = view.map_sync(f, data)
         self.assertEquals(r, map(f, data))
     
+    def test_map_iterable(self):
+        """test map on iterables (direct)"""
+        view = self.client[:]
+        # 101 is prime, so it won't be evenly distributed
+        arr = range(101)
+        # ensure it will be an iterator, even in Python 3
+        it = iter(arr)
+        r = view.map_sync(lambda x:x, arr)
+        self.assertEquals(r, list(arr))
+    
     def test_scatterGatherNonblocking(self):
         data = range(16)
         view = self.client[:]
@@ -446,6 +456,5 @@ class TestView(ClusterTestCase):
                     self.fail(e.evalue)
                 else:
                     raise e
-        
-
+