Abstract

Context variables provide a generic mechanism for tracking dynamic, context-local state, similar to thread-local storage but generalized to cope work with other kinds of thread-like contexts, such as asyncio Tasks. PEP 550 proposed a mechanism for context-local state that was also sensitive to generator context, but this was pretty complicated, so the BDFL requested it be simplified. The result was PEP 567, which is targeted for inclusion in 3.7. This PEP then extends PEP 567’s machinery to add generator context sensitivity.

This PEP is starting out in the “deferred” status, because there isn’t enough time to give it proper consideration before the 3.7 feature freeze. The only goal right now is to understand what would be required to add generator context sensitivity in 3.8, so that we can avoid shipping something in 3.7 that would rule it out by accident. (Ruling it out on purpose can wait until 3.8 ;-).)

Rationale

[Currently the point of this PEP is just to understand how this would work, with discussion of whether it’s a good idea deferred until after the 3.7 feature freeze. So rationale is TBD.]

High-level summary

Instead of holding a single Context, the threadstate now holds a ChainMap of Contexts. ContextVar.get and ContextVar.set are backed by the ChainMap. Generators and async generators each have an associated Context that they push onto the ChainMap while they’re running to isolate their context-local changes from their callers, though this can be overridden in cases like @contextlib.contextmanager where “leaking” context changes from the generator into its caller is desirable.

Specification

class Context(collections.abc.Mapping):
    def __init__(self):
        self._data = {}
        self._in_use = False

    def __getitem__(self, key):
        return self._data[key]

    def __iter__(self):
        return iter(self._data)

    def __len__(self):
        return len(self._data)

# Context.run
def run(self, fn, *args, **kwargs):
    if self._in_use:
        raise RuntimeError("Context already in use")
    tstate = get_thread_state()
    old_context = tstate.current_context
    tstate.current_context = self
    self._in_use = True
    try:
        return fn(*args, **kwargs)
    finally:
        state.current_context = old_context
        self._in_use = False

def copy_context():
    tstate = get_thread_state()
    new_context = Context()
    new_context._data = dict(tstate.current_context)
    return new_context

class Token:
    MISSING = sentinel_value()

    # Note: constructor is private
    def __init__(self, context, var, old_value):
        self._context = context
        self.var = var
        self.old_value = old_value

    # XX: PEP 567 currently makes this a method on ContextVar, but
    # I'm going to propose it switch to this API because it's simpler.
    def reset(self):
        # XX: should we allow token reuse?
        # XX: should we allow tokens to be used if the saved
        # context is no longer active?
        if self.old_value is self.MISSING:
            del self._context._data[self.context_var]
        else:
            self._context._data[self.context_var] = self.old_value

# XX: the handling of defaults here uses the simplified proposal from
# https://mail.python.org/pipermail/python-dev/2018-January/151596.html
# This can be updated to whatever we settle on, it was just less
# typing this way :-)
class ContextVar:
    def __init__(self, name, *, default=None):
        self.name = name
        self.default = default

    def get(self):
        context = get_thread_state().current_context
        return context.get(self, self.default)

    def set(self, new_value):
        context = get_thread_state().current_context
        token = Token(context, self, context.get(self, Token.MISSING))
        context._data[self] = new_value
        return token

# Context.run
def run(self, fn, *args, **kwargs):
    if self._in_use:
        raise RuntimeError("Context already in use")
    tstate = get_thread_state()
    old_context_stack = tstate.current_context_stack
    tstate.current_context_stack = ChainMap([self])     # changed
    self._in_use = True
    try:
        return fn(*args, **kwargs)
    finally:
        state.current_context_stack = old_context_stack
        self._in_use = False

# Context.push
def push(self, fn, *args, **kwargs):
    if self._in_use:
        raise RuntimeError("Context already in use")
    tstate = get_thread_state()
    tstate.current_context_stack.maps.insert(0, self)  # different from run
    self._in_use = True
    try:
        return fn(*args, **kwargs)
    finally:
        tstate.current_context_stack.maps.pop(0)       # different from run
        self._in_use = False

# GeneratorType.__next__
def __next__(self):
    if self.context is not None:
        return self.context.push(self.__real_next__)
    else:
        return self.__real_next__()

# contextlib._GeneratorContextManagerBase.__init__
def __init__(self, func, args, kwds):
    self.gen = func(*args, **kwds)
    self.gen.context = None                  # added
    ...

@contextmanager
def decimal_precision(prec):
    with decimal.localcontext() as ctx:
        ctx.prec = prec
        yield

with decimal_precision(2):
    ...

def copy_context() -> Context:
    tstate = get_thread_state()
    new_context = Context()
    new_context._data = dict(tstate.current_context_stack)
    return new_context

# ContextVar.get
def get(self):
    context_stack = get_thread_state().current_context_stack
    return context_stack.get(self, self.default)

# ContextVar.set
def set(self, new_value):
    top_context = get_thread_state().current_context_stack.maps[0]
    token = Token(top_context, self, top_context.get(self, Token.MISSING))
    top_context._data[self] = new_value
    return token

def get_context_stack() -> List[Context]:
    return list(get_thread_state().current_context_stack.maps)

Comparison to PEP 550

The main difference from PEP 550 is that it reified what we’re calling “contexts” and “context stacks” as two different concrete types (LocalContext and ExecutionContext respectively). This led to lots of confusion about what the differences were, and which object should be used in which places. This proposal simplifies things by only reifying the Context, which is “just a dict”, and makes the “context stack” an unnamed feature of the interpreter’s runtime state – though it is still possible to introspect it using get_context_stack, for debugging and other purposes.

Implementation notes

Context will continue to use a HAMT-based mapping structure under the hood instead of dict, since we expect that calls to copy_context are much more common than ContextVar.set. In almost all cases, copy_context will find that there’s only one Context in the stack (because it’s rare for generators to spawn new tasks), and can simply re-use it directly; in other cases HAMTs are cheap to merge and this can be done lazily.

Rather than using an actual ChainMap object, we’ll represent the context stack using some appropriate structure – the most appropriate options are probably either a bare list with the “top” of the stack being the end of the list so we can use push/pop, or else an intrusive linked list (PyThreadState → Context → Context → …), with the “top” of the stack at the beginning of the list to allow efficient push/pop.

A critical optimization in PEP 567 is the caching of values inside ContextVar. Switching from a single context to a context stack makes this a little bit more complicated, but not too much. Currently, we invalidate the cache whenever the threadstate’s current Context changes (on thread switch, and when entering/exiting Context.run). The simplest approach here would be to invalidate the cache whenever stack changes (on thread switch, when entering/exiting Context.run, and when entering/leaving Context.push). The main effect of this is that iterating a generator will invalidate the cache. It seems unlikely that this will cause serious problems, but if it does, then I think it can be avoided with a cleverer cache key that recognizes that pushing and then popping a Context returns the threadstate to its previous state. (Idea: store the cache key for a particular stack configuration in the topmost Context.)

It seems unavoidable in this design that uncached get will be O(n), where n is the size of the context stack. However, n will generally be very small – it’s roughly the number of nested generators, so usually n=1, and it will be extremely rare to see n greater than, say, 5. At worst, n is bounded by the recursion limit. In addition, we can expect that in most cases of deep generator recursion, most of the Contexts in the stack will be empty, and thus can be skipped extremely quickly during lookup. And for repeated lookups the caching mechanism will kick in. So it’s probably possible to construct some extreme case where this causes performance problems, but ordinary code should be essentially unaffected.

Copyright

This document has been placed in the public domain.