Make it easier to get all pending results.

2 files changed, 16 insertions, 5 deletions

diff --git a/bgproc.py b/bgproc.py
index db148b2..3a1e726 100644
--- a/bgproc.py
+++ b/bgproc.py
@@ -65,26 +65,30 @@ class BackgroundProcessing(object):
         for p in self.processes:
             p.join()
         
-    def get_results(self, block=False):
+    def get_results(self, block=False, block_all=False):
         '''Return currently available results.
         
         If 'block' is True, then wait until there is at least one
-        result.
+        result. If `block_all` is True, then wait until all results are in.
+        
+        If no results are pending, then return an empty list, even if
+        blocking is requested (because it would be an eternal block).
         
         '''
 
-        if block and self.pending == 0:
+        if (block or block_all) and self.pending == 0:
             return []
 
         items = []
+        do_block = block or block_all
         while True:
             try:
-                item = self.results.get(block)
+                item = self.results.get(do_block)
             except Queue.Empty:
                 break
             else:
                 self.pending -= 1
                 items.append(item)
-                block = False
+                do_block = block_all and self.pending > 0
         return items
 
diff --git a/bgproc_tests.py b/bgproc_tests.py
index 92e48cd..154d45c 100644
--- a/bgproc_tests.py
+++ b/bgproc_tests.py
@@ -50,3 +50,10 @@ class BackgroundProcessingTests(unittest.TestCase):
         self.assertEqual(self.bg.pending, 0)
         self.assertEqual(self.bg.get_results(block=True), [])
 
+    def test_gets_all_results(self):
+        self.bg.enqueue_request(0)
+        self.bg.enqueue_request(1)
+        self.assertEqual(sorted(self.bg.get_results(block_all=True)), [1, 2])
+        self.assertEqual(self.bg.pending, 0)
+        self.assertEqual(self.bg.get_results(block_all=True), [])
+