distinguish between cancellation of shielded task and current task

Question

While reading: https://docs.python.org/3/library/asyncio-task.html#asyncio.Task.cancel it seems that catching CancelledError is used for two purposes.

One is potentially preventing your task from getting cancelled.

The other is determining that something has cancelled the task you are awaiting. How to tell the difference?

async def cancel_me():
    try:
        await asyncio.sleep(3600)
    except asyncio.CancelledError:
        raise
    finally:
        print('cancel_me(): after sleep')

async def main():
    task = asyncio.create_task(cancel_me())
    await asyncio.sleep(1)
    task.cancel()
    try:
        await task
    except asyncio.CancelledError:
        # HERE: How do I know if `task` has been cancelled, or I AM being cancelled?
        print("main(): cancel_me is cancelled now")

Answer 1

How to tell the difference [between ourselves getting canceled and the task we're awaiting getting canceled]?

Asyncio doesn't make it easy to tell the difference. When an outer task awaits an inner task, it is delegating control to inner one's coroutine. As a result, canceling either task injects a CancelledError into the exact same place: the inner-most await inside the inner task. This is why you cannot tell the which of the two tasks was canceled originally.

However, it is possible to circumvent the issue by breaking the chain of awaits and connecting the tasks using a completion callback instead. Cancellation of the inner task is then intercepted and detected in the callback:

class ChildCancelled(asyncio.CancelledError):
    pass

async def detect_cancel(task):
    cont = asyncio.get_event_loop().create_future()
    def on_done(_):
        if task.cancelled():
            cont.set_exception(ChildCancelled())
        elif task.exception() is not None:
            cont.set_exception(task.exception())
        else:
            cont.set_result(task.result())
    task.add_done_callback(on_done)
    await cont

This is functionally equivalent to await task, except it doesn't await the inner task directly; it awaits a dummy future whose result is set after task completes. At this point we can replace the CancelledError (which we know must have come from cancellation of the inner task) with the more specific ChildCancelled. On the other hand, if the outer task is cancelled, that will show up as a regular CancelledError at await cont, and will be propagated as usual.

Here is some test code:

import asyncio, sys

# async def detect_cancel defined as above

async def cancel_me():
    print('cancel_me')
    try:
        await asyncio.sleep(3600)
    finally:
        print('cancel_me(): after sleep')

async def parent(task):
    await asyncio.sleep(.001)
    try:
        await detect_cancel(task)
    except ChildCancelled:
        print("parent(): child is cancelled now")
        raise
    except asyncio.CancelledError:
        print("parent(): I am cancelled")
        raise

async def main():
    loop = asyncio.get_event_loop()
    child_task = loop.create_task(cancel_me())
    parent_task = loop.create_task(parent(child_task))
    await asyncio.sleep(.1)  # give a chance to child to start running
    if sys.argv[1] == 'parent':
        parent_task.cancel()
    else:
        child_task.cancel()
    await asyncio.sleep(.5)

asyncio.get_event_loop().run_until_complete(main())

Note that with this implementation, cancelling the outer task will not automatically cancel the inner one, but that may be easily changed with an explicit call to child.cancel(), either in parent, or in detect_cancel itself.

Asyncio uses a similar approach to implement asyncio.shield().