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