How do I write consistent stateful context managers?

Question

EDIT: As pointed out by Thierry Lathuille, PEP567, where ContextVar was introduced, was not designed to address generators (unlike the withdrawn PEP550). Still, the main question remains. How do I write stateful context managers that behave correctly with multiple threads, generators and asyncio tasks?

---

I have a library with some functions that can work in different "modes", so their behavior can be altered by a local context. I am looking at the contextvars module to implement this reliably, so I can use it from different threads, asynchronous contexts, etc. However, I am having trouble getting a simple example working right. Consider this minimal setup:

from contextlib import contextmanager
from contextvars import ContextVar

MODE = ContextVar('mode', default=0)

@contextmanager
def use_mode(mode):
    t = MODE.set(mode)
    try:
        yield
    finally:
        MODE.reset(t)

def print_mode():
   print(f'Mode {MODE.get()}')

Here is a small test with a generator function:

def first():
    print('Start first')
    print_mode()
    with use_mode(1):
        print('In first: with use_mode(1)')
        print('In first: start second')
        it = second()
        next(it)
        print('In first: back from second')
        print_mode()
        print('In first: continue second')
        next(it, None)
        print('In first: finish')

def second():
    print('Start second')
    print_mode()
    with use_mode(2):
        print('In second: with use_mode(2)')
        print('In second: yield')
        yield
        print('In second: continue')
        print_mode()
        print('In second: finish')

first()

I get the following output:

Start first
Mode 0
In first: with use_mode(1)
In first: start second
Start second
Mode 1
In second: with use_mode(2)
In second: yield
In first: back from second
Mode 2
In first: continue second
In second: continue
Mode 2
In second: finish
In first: finish

In the section:

In first: back from second
Mode 2
In first: continue second

It should be Mode 1 instead of Mode 2, because this was printed from first, where the applying context should be, as I understand it, use_mode(1). However, it seems that the use_mode(2) of second is stacked over it until the generator finishes. Are generators not supported by contextvars? If so, is there any way to support stateful context managers reliably? By reliably, I mean it should behave consistently whether I use:

• Multiple threads.
• Generators.
• asyncio

Answer 1

You've actually got an "interlocked context" there - without returning the __exit__ part for the second function it will not restore the context with ContextVars, no matter what.

So, I came up with something here - and the best thing I could think of is a decorator to explicit declare which callables will have their own context - I created a ContextLocal class which works as a namespace, just like thread.local - and attributes in that namespace should behave properly as you expect.

I am finishing the code now - so I had not tested it yet for async or multi-threading, but it should work. If you can help me write a proper test, the solution below could become a Python package in itself.

(I had to resort to injecting an object in generator and co-routines frames locals dictionary in order to clean up the context registry once a generator or co-routine is over - there is PEP 558 formalizing the behavior of locals() for Python 3.8+, and I don't remember now if this injection is allowed - it works up to 3.8 beta 3, though, so I think this usage is valid).

Anyway, here is the code (named as context_wrapper.py):

"""
Super context wrapper -

meant to be simpler to use and work in more scenarios than
Python's contextvars.

Usage:
Create one or more project-wide instances of "ContextLocal"
Decorate your functions, co-routines, worker-methods and generators
that should hold their own states with that instance's `context` method -

and use the instance as namespace for private variables that will be local
and non-local until entering another callable decorated
with `intance.context` - that will create a new, separated scope
visible inside  the decorated callable.


"""

import sys
from functools import wraps

__author__ = "João S. O. Bueno"
__license__ = "LGPL v. 3.0+"

class ContextError(AttributeError):
    pass


class ContextSentinel:
    def __init__(self, registry, key):
        self.registry = registry
        self.key = key

    def __del__(self):
        del self.registry[self.key]


_sentinel = object()


class ContextLocal:

    def __init__(self):
        super().__setattr__("_registry", {})

    def _introspect_registry(self, name=None):

        f = sys._getframe(2)
        while f:
            h = hash(f)
            if h in self._registry and (name is None or name in self._registry[h]):
                return self._registry[h]
            f = f.f_back
        if name:
            raise ContextError(f"{name !r} not defined in any previous context")
        raise ContextError("No previous context set")


    def __getattr__(self, name):
        namespace = self._introspect_registry(name)
        return namespace[name]


    def __setattr__(self, name, value):
        namespace = self._introspect_registry()
        namespace[name] = value


    def __delattr__(self, name):
        namespace = self._introspect_registry(name)
        del namespace[name]

    def context(self, callable_):
        @wraps(callable_)
        def wrapper(*args, **kw):
            f = sys._getframe()
            self._registry[hash(f)] = {}
            result = _sentinel
            try:
                result = callable_(*args, **kw)
            finally:
                del self._registry[hash(f)]
                # Setup context for generator or coroutine if one was returned:
                if result is not _sentinel:
                    frame = getattr(result, "gi_frame", getattr(result, "cr_frame", None))
                    if frame:
                        self._registry[hash(frame)] = {}
                        frame.f_locals["$context_sentinel"] = ContextSentinel(self._registry, hash(frame))

            return result
        return wrapper

Here is the modified version of your example to use with it:

from contextlib import contextmanager

from context_wrapper import ContextLocal

ctx = ContextLocal()


@contextmanager
def use_mode(mode):
    ctx.MODE = mode
    print("entering use_mode")
    print_mode()
    try:
        yield
    finally:

        pass

def print_mode():
   print(f'Mode {ctx.MODE}')


@ctx.context
def first():
    ctx.MODE = 0
    print('Start first')
    print_mode()
    with use_mode(1):
        print('In first: with use_mode(1)')
        print('In first: start second')
        it = second()
        next(it)
        print('In first: back from second')
        print_mode()
        print('In first: continue second')
        next(it, None)
        print('In first: finish')
        print_mode()
    print("at end")
    print_mode()

@ctx.context
def second():
    print('Start second')
    print_mode()
    with use_mode(2):
        print('In second: with use_mode(2)')
        print('In second: yield')
        yield
        print('In second: continue')
        print_mode()
        print('In second: finish')

first()

Here is the output of running that:

Start first
Mode 0
entering use_mode
Mode 1
In first: with use_mode(1)
In first: start second
Start second
Mode 1
entering use_mode
Mode 2
In second: with use_mode(2)
In second: yield
In first: back from second
Mode 1
In first: continue second
In second: continue
Mode 2
In second: finish
In first: finish
Mode 1
at end
Mode 1

(it will be slower than native contextvars by orders of magnitude as those are built-in Python runtime native code - but it seems easier to wrap-the-mind around to use by the same amount)