-

cool-RR · cool-RR · commit 6e4a5a34a0ef · 2015-06-06T19:10:30.000+03:00
diff --git a/source_py3/python_toolbox/future_tools.py b/source_py3/python_toolbox/future_tools.py
@@ -12,14 +12,34 @@ def filter(self, filter_function, iterable, timeout=None,
                as_completed=False):
         '''Get a parallelized version of filter(filter_function, iterable).'''
         
-        sequence = sequence_tools.ensure_iterable_is_sequence(iterable)
-        return (
-            item for (item, keep) in zip(
-                sequence,
-                self.map(filter_function, sequence, timeout=timeout,
-                         as_completed=as_completed)
-            ) if keep
-        )
+        if timeout is not None:
+            end_time = timeout + time.time()
+
+        def make_future(item):
+            future = self.submit(filter_function, item)
+            future._item = item
+            return future
+            
+        futures = tuple(map(make_future, iterable))
+        futures_iterator = concurrent.futures.as_completed(futures) if \
+                                                      as_completed else futures
+
+        # Yield must be hidden in closure so that the futures are submitted
+        # before the first iterator value is required.
+        def result_iterator():
+            try:
+                for future in futures_iterator:
+                    if timeout is None:
+                        result = future.result()
+                    else:
+                        result = future.result(end_time - time.time())
+                    if result:
+                        yield future._item
+            finally:
+                for future in futures:
+                    future.cancel()
+        return result_iterator()
+
 
     def map(self, function, *iterables, timeout=None, as_completed=False):
         '''Get a parallelized version of map(function, iterable).'''
@@ -35,7 +55,7 @@ def map(self, function, *iterables, timeout=None, as_completed=False):
         # before the first iterator value is required.
         def result_iterator():
             try:
-                for future in futures:
+                for future in futures_iterator:
                     if timeout is None:
                         yield future.result()
                     else:
diff --git a/source_py3/test_python_toolbox/test_future_tools/test_future_tools.py b/source_py3/test_python_toolbox/test_future_tools/test_future_tools.py
@@ -14,7 +14,7 @@ def sleep_and_return(seconds):
         return seconds
         
     
-    with future_tools.CuteThreadPoolExecutor(5) as executor:
+    with future_tools.CuteThreadPoolExecutor(10) as executor:
         assert isinstance(executor, future_tools.CuteThreadPoolExecutor)
         assert tuple(executor.filter(lambda x: (x % 2 == 0), range(10))) == \
                                                          tuple(range(0, 10, 2))
@@ -23,7 +23,17 @@ def sleep_and_return(seconds):
                                                           list(range(0, 10, 2))
         assert tuple(executor.filter(
             lambda x: (sleep_and_return(x) % 2 == 0), range(9, -1, -1),
-            as_completed=True)) == (range(0, 10, 2))
+            as_completed=True)) == tuple(range(0, 10, 2))
+        
+        
+        assert tuple(executor.map(lambda x: x % 3, range(10))) == \
+                                                 (0, 1, 2, 0, 1, 2, 0, 1, 2, 0)
+        assert sorted(executor.map(lambda x: x % 3, range(10),
+                                     timeout=10**5, as_completed=True)) == \
+                                                 [0, 0, 0, 0, 1, 1, 1, 2, 2, 2]
+
+        assert tuple(executor.map(sleep_and_return, range(9, -1, -1),
+                                  as_completed=True)) == tuple(range(10))
         
         
         
