是否可以在python中组合异步上下文管理器?与asyncio.gather
类似,但可以与上下文管理器一起使用。像这样:
async def foo():
async with asyncio.gather_cm(start_vm(), start_vm()) as vm1, vm2:
await vm1.do_something()
await vm2.do_something()
这当前可行吗?
答案 0 :(得分:6)
Python 3.7中引入的AsyncExitStack
可以实现接近gather_cm
的地方:
async def foo():
async with AsyncExitStack() as stack:
vm1, vm2 = await asyncio.gather(
stack.enter_async_context(start_vm()),
stack.enter_async_context(start_vm()))
await vm1.do_something()
await vm2.do_something()
不幸的是,__aexit__
仍将顺序运行。这是因为AsyncExitStack
模拟了嵌套的上下文管理器,它们具有明确定义的顺序并且不能重叠。向外部上下文管理器的__aexit__
提供有关内部上下文管理器是否引发异常的信息。 (如果发生异常,数据库句柄的__aexit__
可能会使用它来回滚事务,否则将其提交。)并行运行__aexit__
会使上下文管理器重叠,并且异常信息不可用或不可靠。因此,尽管gather(...)
并行运行__aenter__
,但是AsyncExitStack
记录哪个先出现并以相反的顺序运行__aexit__
。
使用异步上下文管理器,可以使用gather_cm
之类的替代方法。可以放弃嵌套语义,并提供一个聚合上下文管理器,其工作方式类似于“退出池”而不是堆栈。出口池采用了许多彼此独立的上下文管理器,这使得它们的__aenter__
和__aexit__
方法可以并行运行。
棘手的部分是正确处理异常:如果引发任何__aenter__
,则必须传播该异常以防止运行with
块。为了确保正确性,池必须保证__aexit__
已完成的所有上下文管理器都将调用__aenter__
。
这是一个示例实现:
import asyncio
import sys
class gather_cm:
def __init__(self, *cms):
self._cms = cms
async def __aenter__(self):
futs = [asyncio.create_task(cm.__aenter__())
for cm in self._cms]
await asyncio.wait(futs)
# only exit the cms we've successfully entered
self._cms = [cm for cm, fut in zip(self._cms, futs)
if not fut.cancelled() and not fut.exception()]
try:
return tuple(fut.result() for fut in futs)
except:
await self._exit(*sys.exc_info())
raise
async def _exit(self, *args):
# don't use gather() to ensure that we wait for all __aexit__s
# to complete even if one of them raises
done, _pending = await asyncio.wait(
[cm.__aexit__(*args)
for cm in self._cms if cm is not None])
return all(suppress.result() for suppress in done)
async def __aexit__(self, *args):
# Since exits are running in parallel, so they can't see each
# other exceptions. Send exception info from `async with`
# body to all.
return await self._exit(*args)
此测试程序展示了其工作原理:
class test_cm:
def __init__(self, x):
self.x = x
async def __aenter__(self):
print('__aenter__', self.x)
return self.x
async def __aexit__(self, *args):
print('__aexit__', self.x, args)
async def foo():
async with gather_cm(test_cm('foo'), test_cm('bar')) as (cm1, cm2):
print('cm1', cm1)
print('cm2', cm2)
asyncio.run(foo())