Rename a variable and fix type annotations

patch-async-enter-all
overclockworked64 2021-10-17 07:33:37 +02:00
parent c8e7eb8f0b
commit 5f7802dc01
No known key found for this signature in database
GPG Key ID: 0E4A3AD9DB596812
1 changed files with 16 additions and 16 deletions

View File

@ -2,8 +2,8 @@
Async context manager primitives with hard ``trio``-aware semantics Async context manager primitives with hard ``trio``-aware semantics
''' '''
from typing import AsyncContextManager from typing import AsyncContextManager, AsyncGenerator
from typing import TypeVar from typing import TypeVar, Sequence
from contextlib import asynccontextmanager as acm from contextlib import asynccontextmanager as acm
import trio import trio
@ -14,21 +14,19 @@ T = TypeVar("T")
async def _enter_and_wait( async def _enter_and_wait(
mngr: AsyncContextManager[T], mngr: AsyncContextManager[T],
to_yield: dict[int, T], unwrapped: dict[int, T],
all_entered: trio.Event, all_entered: trio.Event,
teardown_trigger: trio.Event, teardown_trigger: trio.Event,
) -> None:
) -> T:
'''Open the async context manager deliver it's value '''Open the async context manager deliver it's value
to this task's spawner and sleep until cancelled. to this task's spawner and sleep until cancelled.
''' '''
async with mngr as value: async with mngr as value:
to_yield[id(mngr)] = value unwrapped[id(mngr)] = value
if all(to_yield.values()): if all(unwrapped.values()):
all_entered.set() all_entered.set()
await teardown_trigger.wait() await teardown_trigger.wait()
@ -36,13 +34,14 @@ async def _enter_and_wait(
@acm @acm
async def async_enter_all( async def async_enter_all(
mngrs: Sequence[AsyncContextManager[T]],
*mngrs: tuple[AsyncContextManager[T]],
teardown_trigger: trio.Event, teardown_trigger: trio.Event,
) -> AsyncGenerator[tuple[T, ...], None]:
) -> tuple[T]: """This async context manager expects a 'teardown_trigger' from the
outside world which will be used internally to gracefully teardown
to_yield = {}.fromkeys(id(mngr) for mngr in mngrs) individual context managers.
"""
unwrapped = {}.fromkeys(id(mngr) for mngr in mngrs)
all_entered = trio.Event() all_entered = trio.Event()
@ -51,11 +50,12 @@ async def async_enter_all(
n.start_soon( n.start_soon(
_enter_and_wait, _enter_and_wait,
mngr, mngr,
to_yield, unwrapped,
all_entered, all_entered,
teardown_trigger, 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(unwrapped.values())