Rename a variable and fix type annotations

246_facepalm_backup
overclockworked64 2021-10-17 07:33:37 +02:00 committed by Tyler Goodlet
parent 6f9229cd09
commit 3130a04c61
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())