49 lines
1.9 KiB
Python
49 lines
1.9 KiB
Python
from types import CoroutineType
|
|
import asyncio
|
|
|
|
# TODO: write a less confusing interface that allows the code to be written more flatly.
|
|
# maybe like: async for done in apply(doit, [tuple_of_args]):
|
|
|
|
def make_pooler(pool_size, finisher=None):
|
|
aws = set()
|
|
async def pooler(item=None):
|
|
nonlocal aws
|
|
finish = item is None
|
|
if not finish:
|
|
if isinstance(item, CoroutineType):
|
|
assert not isinstance(item, asyncio.Task)
|
|
item = asyncio.create_task(item)
|
|
aws.add(item)
|
|
# TODO: don't wait until all completed, just first completed in loop.
|
|
# that way we can handle each done task ASAP.
|
|
condition = asyncio.ALL_COMPLETED if finish else asyncio.FIRST_COMPLETED
|
|
if len(aws) == 0:
|
|
return None
|
|
if finish or len(aws) >= pool_size:
|
|
done, pending = await asyncio.wait(aws, return_when=condition)
|
|
#pending = set(task for task in pending if task in aws) # ???
|
|
ret = None if finisher is None else finisher(done, pending)
|
|
#aws = set(task for task in pending if not task.cancelled())
|
|
aws = pending
|
|
if ret is not None:
|
|
return ret
|
|
return None
|
|
return pooler
|
|
|
|
def make_simple_pooler(pool_size, finisher=None):
|
|
condition = asyncio.FIRST_COMPLETED
|
|
pending = set()
|
|
async def pooler(item=None):
|
|
nonlocal pending
|
|
finish = item is None
|
|
if not finish:
|
|
if isinstance(item, CoroutineType):
|
|
assert not isinstance(item, asyncio.Task)
|
|
item = asyncio.create_task(item)
|
|
pending.add(item)
|
|
desired_size = 0 if finish else pool_size - 1
|
|
while len(pending) > desired_size:
|
|
done, pending = await asyncio.wait(pending, return_when=condition)
|
|
finisher(done, pending)
|
|
return pooler
|