Show how Context Managers do not correctly support suspended execution

async_context.py

import asyncio

class CM:
    mode = 0

    def __init__(self, m):
        self.m = m

    def __enter__(self):
        self.old = CM.mode
        CM.mode = self.m

    def __exit__(self, *_):
        CM.mode = self.old
        

async def long_op():
    with CM('long'):
        print(f'- - mode {CM.mode}')
        await asyncio.sleep(3)
        print(f'- - mode {CM.mode} done')

async def short_op():
    with CM('short'):
        for i in range(5):
            print(f'{i} - mode {CM.mode}')
            await asyncio.sleep(1)
            print(f'{i} - mode {CM.mode} done')
    

async def main():
    await asyncio.gather(short_op(),
                         long_op())

asyncio.run(main())

Note that the short_op iterations have the wrong mode while long_op is waiting.

with CM('short'):
    for i in range(5):
        print(f'{i} - mode {CM.mode}')
        await asyncio.sleep(1)
        print(f'{i} - mode {CM.mode} done')

You would expect this to have mode be 'short' for all iterations, but this is not the case - the value of mode depends on the context manager in long(). Note how the short() iterations have the wrong mode while long() is still waiting:

0 - mode short
- - mode long
0 - mode long done
1 - mode long
1 - mode long done
2 - mode long
- - mode long done
2 - mode short done
3 - mode short
3 - mode short done
4 - mode short
4 - mode short done

---

The expected output is:

0 - mode short
- - mode long
0 - mode short done
1 - mode short
1 - mode short done
2 - mode short
- - mode long done
2 - mode short done
3 - mode short
3 - mode short done
4 - mode short
4 - mode short done

Where each iteration of short_op occurs in the correct context, regardless of the state of long_op. This would be accomplished by extending the definition of CM with __suspend__ and __resume__ methods:

class CM:
    mode = 0
    
    def __init__(self, m):
        self.m = m

    def __enter__(self):
        self.old = CM.mode
        CM.mode = self.m

    def __resume__(self):
        self.old = CM.mode
        CM.mode = self.m

    def __exit__(self, *_):
        CM.mode = self.old

    def __suspend__(self):
        CM.mode = self.old

Or by using a suspendable decorator:

@suspendable
class CM:
    mode = 0
    
    def __init__(self, m):
        self.m = m

    def __enter__(self):
        self.old = CM.mode
        CM.mode = self.m

    def __exit__(self, *_):
        CM.mode = self.old

This is how your example can be fixed with the contextvars module:

import asyncio
import contextvars

cm_mode = contextvars.ContextVar('cm_mode', default=None)

class CM:

    def __init__(self, m):
        self.m = m

    @property
    def mode(self):
        return cm_mode.get()

    def __enter__(self):
        self.old = cm_mode.get()
        cm_mode.set(self.m)
        return self

    def __exit__(self, *_):
        cm_mode.set(self.old)


async def long_op():
    with CM('long') as cm:
        print(f'- - mode {cm.mode}')
        await asyncio.sleep(3)
        print(f'- - mode {cm.mode} done')

async def short_op():
    with CM('short') as cm:
        for i in range(5):
            print(f'{i} - mode {cm.mode}')
            await asyncio.sleep(1)
            print(f'{i} - mode {cm.mode} done')


async def main():
    await asyncio.gather(short_op(),
                         long_op())

asyncio.run(main())