Make 'async_enter_all' take a teardown trigger which '_enter_and_wait' will wait on
parent
21afc69ac7
commit
2815f1c343
|
@ -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()
|
|
||||||
|
|
Loading…
Reference in New Issue