From 65a645bdde31be8df6406047bc8f2d3e89aeb64e Mon Sep 17 00:00:00 2001 From: Tyler Goodlet Date: Fri, 1 Oct 2021 10:27:08 -0400 Subject: [PATCH] Start trionics mod with an `async_enter_all` --- piker/trionics.py | 80 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 80 insertions(+) create mode 100644 piker/trionics.py diff --git a/piker/trionics.py b/piker/trionics.py new file mode 100644 index 00000000..10f6a33d --- /dev/null +++ b/piker/trionics.py @@ -0,0 +1,80 @@ +# piker: trading gear for hackers +# Copyright (C) Tyler Goodlet (in stewardship of piker0) + +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. + +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. + +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +''' +sugarz for trio/tractor conc peeps. + +''' +from typing import AsyncContextManager +from typing import TypeVar +from contextlib import asynccontextmanager as acm + +import trio + + +# A regular invariant generic type +T = TypeVar("T") + + +async def _enter_and_sleep( + + mngr: AsyncContextManager[T], + to_yield: dict[int, T], + all_entered: trio.Event, + # task_status: TaskStatus[T] = trio.TASK_STATUS_IGNORED, + +) -> T: + '''Open the async context manager deliver it's value + to this task's spawner and sleep until cancelled. + + ''' + async with mngr as value: + to_yield[id(mngr)] = value + + if all(to_yield.values()): + all_entered.set() + + # sleep until cancelled + await trio.sleep_forever() + + +@acm +async def async_enter_all( + + *mngrs: list[AsyncContextManager[T]], + +) -> tuple[T]: + + to_yield = {}.fromkeys(id(mngr) for mngr in mngrs) + + all_entered = trio.Event() + + async with trio.open_nursery() as n: + for mngr in mngrs: + n.start_soon( + _enter_and_sleep, + mngr, + to_yield, + all_entered, + ) + + # 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()