Determine whether task was cancelled from inside or outside

distinguish_cancellation.py

import asyncio


class CancelledFromOutside(asyncio.CancelledError):
    pass


class CancelledFromInside(asyncio.CancelledError):
    pass


async def distinguish_cancellation(fut):
    """Wait for a future. If cancelled, raise different exceptions depending
    on who did the cancellation.

    If fut was cancelled, propagate cancellation outward by raising
    CancelledFromInside.
    If this function was cancelled, cancel fut, and raise CancelledFromOutside.
    """
    fut = asyncio.ensure_future(fut)
    try:
        await asyncio.wait((fut,))
    except asyncio.CancelledError:
        fut.cancel()
        raise CancelledFromOutside
    assert fut.done()
    if fut.cancelled():
        raise CancelledFromInside
    return fut.result()