diff options
-rw-r--r-- | bgproc.py | 14 | ||||
-rw-r--r-- | bgproc_tests.py | 7 |
2 files changed, 16 insertions, 5 deletions
@@ -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), []) + |