Make 'async_enter_all' take a teardown trigger which '_enter_and_wait' will wait on

patch-async-enter-all
overclockworked64 2021-10-16 17:34:59 +02:00
parent 21042ef926
commit dc8ccca8be
No known key found for this signature in database
GPG Key ID: 0E4A3AD9DB596812
1 changed files with 6 additions and 9 deletions

View File

@ -13,12 +13,12 @@ import trio
T = TypeVar("T") T = TypeVar("T")
async def _enter_and_sleep( async def _enter_and_wait(
mngr: AsyncContextManager[T], mngr: AsyncContextManager[T],
to_yield: dict[int, T], to_yield: dict[int, T],
all_entered: trio.Event, all_entered: trio.Event,
# task_status: TaskStatus[T] = trio.TASK_STATUS_IGNORED, teardown_trigger: trio.Event,
) -> T: ) -> T:
'''Open the async context manager deliver it's value '''Open the async context manager deliver it's value
@ -31,14 +31,14 @@ async def _enter_and_sleep(
if all(to_yield.values()): if all(to_yield.values()):
all_entered.set() all_entered.set()
# sleep until cancelled await teardown_trigger.wait()
await trio.sleep_forever()
@acm @acm
async def async_enter_all( async def async_enter_all(
*mngrs: tuple[AsyncContextManager[T]], *mngrs: tuple[AsyncContextManager[T]],
teardown_trigger: trio.Event,
) -> tuple[T]: ) -> tuple[T]:
@ -49,16 +49,13 @@ async def async_enter_all(
async with trio.open_nursery() as n: async with trio.open_nursery() as n:
for mngr in mngrs: for mngr in mngrs:
n.start_soon( n.start_soon(
_enter_and_sleep, _enter_and_wait,
mngr, mngr,
to_yield, to_yield,
all_entered, all_entered,
teardown_trigger,
) )
# deliver control once all managers have started up # deliver control once all managers have started up
await all_entered.wait() await all_entered.wait()
yield tuple(to_yield.values()) yield tuple(to_yield.values())
# tear down all sleeper tasks thus triggering individual
# mngr ``__aexit__()``s.
n.cancel_scope.cancel()