Open graingert opened 2 years ago
I think Cancelled needs special handling:
def aio_acmgr_fn_as_trio_acmgr_fn(acmgr_fn):
@functools.wraps(acmgr_fn)
@contextlib.asynccontextmanager
async def wrapper(*args, **kwargs):
@trio_asyncio.aio_as_trio
async def run(*, task_status):
fut = asyncio.get_running_loop().create_future()
async with acmgr_fn(*args, **kwargs) as result:
task_status.started((fut, result))
await fut
async with trio.open_nursery() as nursery:
fut, result = await nursery.start(run)
try:
yield result
except trio.Cancelled:
raise
except BaseException as e:
fut.set_exception(e)
finally:
if not fut.done():
fut.set_result(None)
return wrapper
currently @aio_as_trio runs
__aenter__
and__aexit__
in two different tasks, whereas some asyncio context managers are sensitive to being run in the same taskso I've got a little util function like this:
I think it would be worth including it here also