Skip to content

Commit

Permalink
Ensure user errors cleanup internal futures
Browse files Browse the repository at this point in the history
  • Loading branch information
mikeshardmind committed Jan 15, 2025
1 parent f43887e commit 01e4891
Showing 1 changed file with 13 additions and 8 deletions.
21 changes: 13 additions & 8 deletions src/async_utils/gen_transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,16 +97,21 @@ def sync_to_async_gen[**P, Y](
background_task = asyncio.create_task(background_coro)

async def gen() -> AsyncGenerator[Y]:
q_get = None
try:
while not background_task.done():
q_get = asyncio.ensure_future(q.async_get())
done, _pending = await asyncio.wait(
(background_task, q_get),
return_when=asyncio.FIRST_COMPLETED,
)
if q_get in done:
yield (await q_get)
laziness_ev.set()
try:
q_get = asyncio.ensure_future(q.async_get())
done, _pending = await asyncio.wait(
(background_task, q_get),
return_when=asyncio.FIRST_COMPLETED,
)
if q_get in done:
yield (await q_get)
laziness_ev.set()
finally:
if q_get is not None:
q_get.cancel()
laziness_ev.clear()
while q:
yield (await q.async_get())
Expand Down

0 comments on commit 01e4891

Please sign in to comment.