Skip to content

Commit

Permalink
use wait instead of as_completed
Browse files Browse the repository at this point in the history
  • Loading branch information
ebonnal committed Sep 28, 2024
1 parent 585cbaf commit b8fa908
Showing 1 changed file with 14 additions and 7 deletions.
21 changes: 14 additions & 7 deletions streamable/iters.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,13 @@
from abc import ABC, abstractmethod
from asyncio import AbstractEventLoop, get_event_loop
from collections import defaultdict, deque
from concurrent.futures import Executor, Future, ThreadPoolExecutor, as_completed
from concurrent.futures import (
FIRST_COMPLETED,
Executor,
Future,
ThreadPoolExecutor,
wait,
)
from contextlib import contextmanager
from datetime import datetime
from math import ceil
Expand Down Expand Up @@ -424,9 +430,10 @@ def _get_result(
if self.ordered:
return futures.popleft().result()
else:
future = next(as_completed(futures))
futures.remove(future)
return future.result()
done_futures, _ = wait(futures, return_when=FIRST_COMPLETED)
done_future = next(iter(done_futures))
futures.remove(done_future)
return done_future.result()


class AsyncConcurrentMappingIterable(ConcurrentMappingIterable[T, U]):
Expand Down Expand Up @@ -476,9 +483,9 @@ def _get_result(
done_futures, _ = self.loop.run_until_complete(
asyncio.wait(futures, return_when=asyncio.FIRST_COMPLETED) # type: ignore
)
future = next(iter(done_futures))
futures.remove(future)
return future.result()
done_future = next(iter(done_futures))
futures.remove(done_future)
return done_future.result()


class ConcurrentFlatteningIterable(
Expand Down

0 comments on commit b8fa908

Please sign in to comment.