How to write a context manager that handles global state properly

Python Help
1

How do you write a context that manages global state properly? I tried:

import contextlib

_foo_context = None

def get_foo()->str|None:
    return _foo_context

@contextlib.contextmanager
def with_foo(new_context:str):
    global _foo_context
    _foo_context, old_context = new_context, _foo_context
    yield
    _foo_context = old_context

But that fails when used with generators(and async):

def cook_food(item:str):
    with with_foo(item):
        print(f"Cooking {get_foo()}")
        yield
        print(f"Still cooking {get_foo()}")

raw_eggs = cook_food("eggs")
raw_spam = cook_food("spam")
next(raw_eggs,None)
next(raw_spam,None)
next(raw_eggs,None) # prints "Still cooking spam", instead of eggs.
next(raw_spam,None) # prints "Still cooking None"!

How do I do this properly?

2

Do you know why it’s doing that?

When it calls the first next(raw_eggs,None), it resumes cook_food with “eggs”, setting _foo_context to “eggs”, and then yields.

When it calls the first next(raw_spam,None), it resumes cook_food with “spam”, setting _foo_context to “spam”, and then yields.

When it calls the second next(raw_eggs,None), it resumes the first cook_food, printing “spam” because _foo_context is still set to “spam”.

You can’t have 2 context managers sharing a global variable, yet storing 2 separate values in it, at the same time.

2 Likes
3

There is a way to manage global state that works with async and threads, and that is to use contextvars instead of global variables. For example:

import asyncio
import contextlib
from contextvars import ContextVar

_foo_context = ContextVar("_foo_context", default=None)

def get_foo()->str|None:
    return _foo_context.get()

@contextlib.contextmanager
def with_foo(new_context:str):
    token = _foo_context.set(new_context)
    try:
        yield
    finally:
        _foo_context.reset(token)

async def cook_food(item:str):
    with with_foo(item):
        print(f"Cooking {get_foo()}")
        await asyncio.sleep(0.1)
        print(f"Still cooking {get_foo()}")


async def main():
    tasks = [
        asyncio.create_task(cook_food("eggs")),
        asyncio.create_task(cook_food("spam")),
    ]
    await asyncio.wait(tasks)


if __name__ == "__main__":
    asyncio.run(main())

This prints

Cooking eggs
Cooking spam
Still cooking eggs
Still cooking spam

Unfortunately, it does seem contextvars do not currently work with generators. Changing your original example to use a contextvar instead of a global variable does not change the result.

This was surprising to me, TBH. Doing a quick google search I found that there is already PEP about this. It seems this is a know problem but there is (was) no obvious solution available when contextvars were introduced in Python 3.7: