深入了解 Python 中的 asyncio 库 -shield 函数

shield

asyncio.shield，可用于屏蔽取消操作。到目前为止，我们还没有看到任务取消。

看一个例子:

In : loop = asyncio.get_event_loop()
In : task1 = loop.create_task(a())
In : task2 = loop.create_task(b())
In : task1.cancel()
Out: True
In : await asyncio.gather(task1, task2)
Suspending a
Suspending b
---------------------------------------------------------------------------
CancelledError                            Traceback (most recent call last)
cell_name in async-def-wrapper()
CancelledError:

上例中，task1取消后，使用asyncio.gather收集结果，直接抛出CancelledError。这里有一个细节，gather 支持 return_exceptions 参数：

In : await asyncio.gather(task1, task2, return_exceptions=True)
Out: [concurrent.futures._base.CancelledError(), 'B']

可以看到task2还是会被执行，但是task1的返回值是一个CancelledError错误，也就是任务已经被取消了。如果你不想在创建后被任何情况取消，你可以使用 asyncio.shield 来保护任务无法成功完成。但要注意一个陷阱，先看错字：

In : task1 = asyncio.shield(a())
In : task2 = loop.create_task(b())
In : task1.cancel()
Out: True
In : await asyncio.gather(task1, task2, return_exceptions=True)
Suspending a
Suspending b
Resuming b
Out: [concurrent.futures._base.CancelledError(), 'B']

可以看到还是有CancelledError错误，协程a还没有执行。正确用法如下：

In : task1 = asyncio.shield(a())
In : task2 = loop.create_task(b())
In : ts = asyncio.gather(task1, task2, return_exceptions=True)
In : task1.cancel()
Out: True
In : await ts
Suspending a
Suspending b
Resuming a
Resuming b
Out: [concurrent.futures._base.CancelledError(), 'B']

可以看出，虽然结果是CancelledError错误，但是看输出可以确认协程确实执行了。所以正确的步骤是：

先创建 GatheringFuture 对象 ts

取消任务

await ts

asynccontextmanager

如果你了解 Python，你可能之前听说过或使用过 contextmanager，一个上下文管理器。通过时序示例了解它的作用：

from contextlib import contextmanager
async def a():
    await asyncio.sleep(3)
    return 'A'
async def b():
    await asyncio.sleep(1)
    return 'B'
async def s1():
    return await asyncio.gather(a(), b())
@contextmanager
def timed(func):
    start = time.perf_counter()
    yield asyncio.run(func())
    print(f'Cost: {time.perf_counter() - start}')

定时函数使用 contextmanager 装饰器来生成协程的运行结果，并计算执行后的时间消耗：

In : from contextmanager import *
In : with timed(s1) as rv:
...:     print(f'Result: {rv}')
...:
Result: ['A', 'B']
Cost: 3.0052654459999992

让我们先体验一下。在 Python 3.7 中增加了 asynccontextmanager，它是 contextmanager 的异步版本，适用于异步函数的执行。上面的例子可以改成如下：

@asynccontextmanager
async def async_timed(func):
    start = time.perf_counter()
    yield await func()
    print(f'Cost: {time.perf_counter() - start}')
async def main():
    async with async_timed(s1) as rv:
        print(f'Result: {rv}')
In : asyncio.run(main())
Result: ['A', 'B']
Cost: 3.00414147500004

async版本的with要用async with，另外要注意yield await func()这句，相当于yield + await func()

PS: contextmanager 和 asynccontextmanager 最好的理解方法是去看源码注释

下一节：