Skip to content
Snippets Groups Projects
Commit dff252da authored by Piotr Maślanka's avatar Piotr Maślanka
Browse files

add done callback

parent 202fefd1
No related branches found
No related tags found
No related merge requests found
...@@ -32,6 +32,22 @@ class FutureCollection: ...@@ -32,6 +32,22 @@ class FutureCollection:
fc.append(other) fc.append(other)
return FutureCollection(fc) return FutureCollection(fc)
def add_done_callback(self, callback, only_one: bool = False) -> None:
"""
Add a callback to a Future to be called on it's completion.
By default, this will add the callback to all futures.
:param callback: callback that takes the completed Future as argument
:param only_one: callback will be added only to a single Future. False by default
:raises IndexError: only_one was given and no Futures in collection!
"""
if only_one:
self.futures[0].add_done_callback(callback)
else:
for future in self.futures:
future.add_done_callback(callback)
def set_running_or_notify_cancel(self) -> bool: def set_running_or_notify_cancel(self) -> bool:
""" """
Call :code:`set_running_or_notify_cancel` on the futures Call :code:`set_running_or_notify_cancel` on the futures
......
...@@ -19,6 +19,32 @@ from satella.exceptions import WouldWaitMore, AlreadyAllocated, Empty ...@@ -19,6 +19,32 @@ from satella.exceptions import WouldWaitMore, AlreadyAllocated, Empty
class TestConcurrent(unittest.TestCase): class TestConcurrent(unittest.TestCase):
def test_future_collection_callbacks_one(self):
a = {'count': 0}
def callback(fut):
nonlocal a
a['count'] += 1
fc = FutureCollection([PythonFuture(), PythonFuture()])
fc.add_done_callback(callback, True)
fc.set_running_or_notify_cancel()
fc.set_result(None)
self.assertEqual(a['count'], 1)
def test_future_collection_callbacks(self):
a = {'count': 0}
def callback(fut):
nonlocal a
a['count'] += 1
fc = FutureCollection([PythonFuture(), PythonFuture()])
fc.add_done_callback(callback)
fc.set_running_or_notify_cancel()
fc.set_result(None)
self.assertEqual(a['count'], 2)
def test_future_collection_cancel(self): def test_future_collection_cancel(self):
fc = FutureCollection([PythonFuture(), PythonFuture()]) fc = FutureCollection([PythonFuture(), PythonFuture()])
self.assertTrue(fc.set_running_or_notify_cancel()) self.assertTrue(fc.set_running_or_notify_cancel())
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment