diff options
Diffstat (limited to 'smart_future.py')
| -rw-r--r-- | smart_future.py | 15 |
1 files changed, 11 insertions, 4 deletions
diff --git a/smart_future.py b/smart_future.py index 2f3cbd9..604c149 100644 --- a/smart_future.py +++ b/smart_future.py @@ -28,9 +28,11 @@ def wait_any( real_futures = [] smart_future_by_real_future = {} completed_futures = set() - for _ in futures: - real_futures.append(_.wrapped_future) - smart_future_by_real_future[_.wrapped_future] = _ + for f in futures: + assert type(f) == SmartFuture + real_futures.append(f.wrapped_future) + smart_future_by_real_future[f.wrapped_future] = f + while len(completed_futures) != len(real_futures): newly_completed_futures = concurrent.futures.as_completed(real_futures) for f in newly_completed_futures: @@ -56,7 +58,11 @@ def wait_all( *, log_exceptions: bool = True, ) -> None: - real_futures = [x.wrapped_future for x in futures] + real_futures = [] + for f in futures: + assert type(f) == SmartFuture + real_futures.append(f.wrapped_future) + (done, not_done) = concurrent.futures.wait( real_futures, timeout=None, return_when=concurrent.futures.ALL_COMPLETED ) @@ -83,6 +89,7 @@ class SmartFuture(DeferredOperand): """ def __init__(self, wrapped_future: fut.Future) -> None: + assert type(wrapped_future) == fut.Future self.wrapped_future = wrapped_future self.id = id_generator.get("smart_future_id") |