python · rhansen · Sep 28, 2024 · Sep 28, 2024 · Sep 28, 2024 · Sep 28, 2024
@@ -123,33 +123,22 @@ Context object management functions:
 
    Enumeration of possible context object watcher events:
 
-   - ``Py_CONTEXT_EVENT_ENTER``: A context has been entered, causing the
-     :term:`current context` to switch to it.  The object passed to the watch
-     callback is the now-current :class:`contextvars.Context` object.  Each
-     enter event will eventually have a corresponding exit event for the same
-     context object after any subsequently entered contexts have themselves been
-     exited.
-   - ``Py_CONTEXT_EVENT_EXIT``: A context is about to be exited, which will
-     cause the :term:`current context` to switch back to what it was before the
-     context was entered.  The object passed to the watch callback is the
-     still-current :class:`contextvars.Context` object.
+   - ``Py_CONTEXT_SWITCHED``: The :term:`current context` has switched to a
+     different context.  The object passed to the watch callback is the
+     now-current :class:`contextvars.Context` object, or None if no context is
+     current.
 
    .. versionadded:: 3.14
 
-.. c:type:: int (*PyContext_WatchCallback)(PyContextEvent event, PyContext* ctx)
+.. c:type:: void (*PyContext_WatchCallback)(PyContextEvent event, PyObject* obj)
 
    Context object watcher callback function.  The object passed to the callback
    is event-specific; see :c:type:`PyContextEvent` for details.
 
-   If the callback returns with an exception set, it must return ``-1``; this
-   exception will be printed as an unraisable exception using
-   :c:func:`PyErr_FormatUnraisable`. Otherwise it should return ``0``.
+   Any pending exception is cleared before the callback is called and restored
+   after the callback returns.
 
-   There may already be a pending exception set on entry to the callback. In
-   this case, the callback should return ``0`` with the same exception still
-   set. This means the callback may not call any other API that can set an
-   exception unless it saves and clears the exception state first, and restores
-   it before returning.
+   If the callback raises an exception it will be ignored.
 
    .. versionadded:: 3.14
 

@@ -150,20 +150,24 @@ Manual Context Management
    considered to be *entered*.
 
    *Entering* a context, which can be done by calling its :meth:`~Context.run`
-   method, makes the context the current context by pushing it onto the top of
-   the current thread's context stack.
+   method or by using it as a :term:`context manager`, makes the context the
+   current context by pushing it onto the top of the current thread's context
+   stack.
 
    *Exiting* from the current context, which can be done by returning from the
-   callback passed to the :meth:`~Context.run` method, restores the current
-   context to what it was before the context was entered by popping the context
-   off the top of the context stack.
+   callback passed to :meth:`~Context.run` or by exiting the :keyword:`with`
+   statement suite, restores the current context to what it was before the
+   context was entered by popping the context off the top of the context stack.
 
    Since each thread has its own context stack, :class:`ContextVar` objects
    behave in a similar fashion to :func:`threading.local` when values are
    assigned in different threads.
 
-   Attempting to enter an already entered context, including contexts entered in
-   other threads, raises a :exc:`RuntimeError`.
+   Attempting to do either of the following raises a :exc:`RuntimeError`:
+
+   * Entering an already entered context, including contexts entered in
+     other threads.
+   * Exiting from a context that is not the current context.
 
    After exiting a context, it can later be re-entered (from any thread).
 
@@ -176,6 +180,50 @@ Manual Context Management
 
    Context implements the :class:`collections.abc.Mapping` interface.
 
+   .. versionadded:: 3.14
+      Added support for the :term:`context management protocol`.
+
+   When used as a :term:`context manager`, the value bound to the identifier
+   given in the :keyword:`with` statement's :keyword:`!as` clause (if present)
+   is the :class:`!Context` object itself.
+
+   Example:
+
+   .. testcode::
+
+      import contextvars
+
+      var = contextvars.ContextVar("var")
+      var.set("initial")
+      print(var.get())  # 'initial'
+
+      # Copy the current Context and enter the copy.
+      with contextvars.copy_context() as ctx:
+          var.set("updated")
+          print(var in ctx)  # 'True'
+          print(ctx[var])  # 'updated'
+          print(var.get())  # 'updated'
+
+      # Exited ctx, so the observed value of var has reverted.
+      print(var.get())  # 'initial'
+      # But the updated value is still recorded in ctx.
+      print(ctx[var])  # 'updated'
+
+      # Re-entering ctx restores the updated value of var.
+      with ctx:
+          print(var.get())  # 'updated'
+
+   .. testoutput::
+      :hide:
+
+      initial
+      True
+      updated
+      updated
+      initial
+      updated
+      updated
+
    .. method:: run(callable, *args, **kwargs)
 
       Enters the Context, executes ``callable(*args, **kwargs)``, then exits the

diff --git a/Doc/whatsnew/3.14.rst b/Doc/whatsnew/3.14.rst
@@ -226,6 +226,13 @@ ast
   (Contributed by Tomas R in :gh:`116022`.)
 
 
+contextvars
+-----------
+
+* Added support for the :term:`context management protocol` to
+  :class:`contextvars.Context`.  (Contributed by Richard Hansen in :gh:`99634`.)
+
+
 ctypes
 ------
 

@@ -29,30 +29,23 @@ PyAPI_FUNC(int) PyContext_Exit(PyObject *);
 
 typedef enum {
     /*
-     * A context has been entered, causing the "current context" to switch to
-     * it.  The object passed to the watch callback is the now-current
-     * contextvars.Context object.  Each enter event will eventually have a
-     * corresponding exit event for the same context object after any
-     * subsequently entered contexts have themselves been exited.
+     * The current context has switched to a different context.  The object
+     * passed to the watch callback is the now-current contextvars.Context
+     * object, or None if no context is current.
      */
-    Py_CONTEXT_EVENT_ENTER,
-    /*
-     * A context is about to be exited, which will cause the "current context"
-     * to switch back to what it was before the context was entered.  The
-     * object passed to the watch callback is the still-current
-     * contextvars.Context object.
-     */
-    Py_CONTEXT_EVENT_EXIT,
+    Py_CONTEXT_SWITCHED = 1,
 } PyContextEvent;
 
 /*
  * Context object watcher callback function.  The object passed to the callback
  * is event-specific; see PyContextEvent for details.
  *
- * if the callback returns with an exception set, it must return -1. Otherwise
- * it should return 0
+ * Any pending exception is cleared before the callback is called and restored
+ * after the callback returns.
+ *
+ * If the callback raises an exception it will be ignored.
  */
-typedef int (*PyContext_WatchCallback)(PyContextEvent, PyContext *);
+typedef void (*PyContext_WatchCallback)(PyContextEvent, PyObject *);
 
 /*
  * Register a per-interpreter callback that will be invoked for context object

@@ -164,7 +164,6 @@ struct _ts {
     PyObject *async_gen_firstiter;
     PyObject *async_gen_finalizer;
 
-    PyObject *context;
     uint64_t context_ver;
 
     /* Unique thread state id. */

@@ -5,7 +5,11 @@
 #  error "this header requires Py_BUILD_CORE define"
 #endif
 
+#include "cpython/context.h"
+#include "cpython/genobject.h"    // PyGenObject
+#include "pycore_genobject.h"     // PyGenObject
 #include "pycore_hamt.h"          // PyHamtObject
+#include "pycore_tstate.h"        // _PyThreadStateImpl
 
 #define CONTEXT_MAX_WATCHERS 8
 
@@ -24,12 +28,81 @@ typedef struct {
 
 struct _pycontextobject {
     PyObject_HEAD
-    PyContext *ctx_prev;
+    PyObject *ctx_prev;
     PyHamtObject *ctx_vars;
     PyObject *ctx_weakreflist;
     int ctx_entered;
 };
 
+// Resets a coroutine's independent context stack to ctx.  If ctx is NULL or
+// Py_None, the coroutine will be a dependent coroutine (its context stack will
+// be empty) upon successful return.  Otherwise, the coroutine will be an
+// independent coroutine upon successful return, with ctx as the sole item on
+// its context stack.
+//
+// The coroutine's existing stack must be empty (NULL) or contain only a single
+// entry (from a previous call to this function).  If the coroutine is
+// currently executing, this function must be called from the coroutine's
+// thread.
+//
+// Unless ctx already equals the coroutine's existing context stack, the
+// context on the existing stack (if one exists) is immediately exited and ctx
+// (if non-NULL) is immediately entered.
+int _PyGen_ResetContext(PyThreadState *ts, PyGenObject *self, PyObject *ctx);
+
+void _PyGen_ActivateContextImpl(_PyThreadStateImpl *tsi, PyGenObject *self);
+void _PyGen_DeactivateContextImpl(_PyThreadStateImpl *tsi, PyGenObject *self);
+
+// Makes the given coroutine's context stack the active context stack for the
+// thread, shadowing (temporarily deactivating) the thread's previously active
+// context stack.  The context stack remains active until deactivated with a
+// call to _PyGen_DeactivateContext, as long as it is not shadowed by another
+// activated context stack.
+//
+// Each activated context stack must eventually be deactivated by calling
+// _PyGen_DeactivateContext.  The same context stack cannot be activated again
+// until deactivated.
+//
+// If the coroutine's context stack is empty this function has no effect.
+//
+// This is called each time a value is sent to a coroutine, so it is inlined to
+// avoid function call overhead in the common case of dependent coroutines.
+static inline void
+_PyGen_ActivateContext(PyThreadState *ts, PyGenObject *self)
+{
+    _PyThreadStateImpl *tsi = (_PyThreadStateImpl *)ts;
+    assert(self->_ctx_chain.prev == NULL);
+    if (self->_ctx_chain.ctx == NULL) {
+        return;
+    }
+    _PyGen_ActivateContextImpl(tsi, self);
+}
+
+// Deactivates the given coroutine's context stack, un-shadowing (reactivating)
+// the thread's previously active context stack.  Does not affect any contexts
+// in the coroutine's context stack (they remain entered).
+//
+// Must not be called if a different context stack is currently shadowing the
+// coroutine's context stack.
+//
+// If the coroutine's context stack is not the active context stack this
+// function has no effect.
+//
+// This is called each time a value is sent to a coroutine, so it is inlined to
+// avoid function call overhead in the common case of dependent coroutines.
+static inline void
+_PyGen_DeactivateContext(PyThreadState *ts, PyGenObject *self)
+{
+    _PyThreadStateImpl *tsi = (_PyThreadStateImpl *)ts;
+    if (tsi->_ctx_chain.prev != &self->_ctx_chain) {
+        assert(self->_ctx_chain.ctx == NULL);
+        assert(self->_ctx_chain.prev == NULL);
+        return;
+    }
+    assert(self->_ctx_chain.ctx != NULL);
+    _PyGen_DeactivateContextImpl(tsi, self);
+}
+
 
 struct _pycontextvarobject {
     PyObject_HEAD

@@ -0,0 +1,79 @@
+#ifndef Py_INTERNAL_CONTEXTCHAIN_H
+#define Py_INTERNAL_CONTEXTCHAIN_H
+
+#ifndef Py_BUILD_CORE
+#  error "this header requires Py_BUILD_CORE define"
+#endif
+
+#include "pytypedefs.h"  // PyObject
+
+
+// Circularly linked chain of multiple independent context stacks, used to give
+// coroutines (including generators) their own (optional) independent context
+// stacks.
+//
+// Detailed notes on how this chain is used:
+//   * The chain is circular simply to save a pointer's worth of memory in
+//     _PyThreadStateImpl.  It is actually used as an ordinary linear linked
+//     list.  It is called "chain" instead of "stack" or "list" to evoke "call
+//     chain", which it is related to, and to avoid confusion with "context
+//     stack".
+//   * There is one chain per thread, and _PyThreadStateImpl::_ctx_chain::prev
+//     points to the head of the thread's chain.
+//   * A thread's chain is never empty.
+//   * _PyThreadStateImpl::_ctx_chain is always the tail entry of the thread's
+//     chain.
+//   * _PyThreadStateImpl::_ctx_chain is usually the only link in the thread's
+//     chain, so _PyThreadStateImpl::_ctx_chain::prev usually points to the
+//     _PyThreadStateImpl::_ctx_chain itself.
+//   * The "active context stack" is always at the head link in a thread's
+//     context chain.  Contexts are entered by pushing onto the active context
+//     stack and exited by popping off of the active context stack.
+//   * The "current context" is the top context in the active context stack.
+//     Context variable accesses (reads/writes) use the current context.
+//   * A *dependent* coroutine or generator is a coroutine or generator that
+//     does not have its own independent context stack.  When a dependent
+//     coroutine starts or resumes execution, the current context -- as
+//     observed by the coroutine -- is the same context that was current just
+//     before the coroutine's `send` method was called.  This means that the
+//     current context as observed by a dependent coroutine can change
+//     arbitrarily during a yield/await.  Dependent coroutines are so-named
+//     because they depend on their senders to enter the appropriate context
+//     before each send.  Coroutines and generators are dependent by default
+//     for backwards compatibility.
+//   * The purpose of the context chain is to enable *independent* coroutines
+//     and generators, which have their own context stacks.  Whenever an
+//     independent coroutine starts or resumes execution, the current context
+//     automatically switches to the context associated with the coroutine.
+//     This is accomplished by linking the coroutine's chain link (at
+//     PyGenObject::_ctx_chain) to the head of the thread's chain.  Independent
+//     coroutines are so-named because they do not depend on their senders to
+//     enter the appropriate context before each send.
+//   * The head link is unlinked from the thread's chain when its associated
+//     independent coroutine or generator stops executing (yields, awaits,
+//     returns, or throws).
+//   * A running dependent coroutine's chain link is linked into the thread's
+//     chain if the coroutine is upgraded from dependent to independent by
+//     assigning a context to the coroutine's `_context` property.  The chain
+//     link is inserted at the position corresponding to the coroutine's
+//     position in the call chain relative to any other currently running
+//     independent coroutines.  For example, if dependent coroutine `coro_a`
+//     calls function `func_b` which resumes independent coroutine `coro_c`
+//     which assigns a context to `coro_a._context`, then `coro_a` becomes an
+//     independent coroutine with its chain link inserted after `coro_c`'s
+//     chain link (which remains the head link).
+//   * A running independent coroutine's chain link is unlinked from the
+//     thread's chain if the coroutine is downgraded from independent to
+//     dependent by assigning `None` to its `_context` property.
+//   * The references to the object at the `prev` link in the chain are
+//     implicit (borrowed).
+typedef struct _PyContextChain {
+    // NULL for dependent coroutines/generators, non-NULL for independent
+    // coroutines/generators.
+    PyObject *ctx;
+    // NULL if unlinked from the thread's context chain, non-NULL otherwise.
+    struct _PyContextChain *prev;
+} _PyContextChain;
+
+
+#endif /* !Py_INTERNAL_CONTEXTCHAIN_H */
@@ -1,5 +1,8 @@
 #ifndef Py_INTERNAL_GENOBJECT_H
 #define Py_INTERNAL_GENOBJECT_H
+
+#include "pycore_contextchain.h"  // _PyContextChain
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -22,6 +25,7 @@ extern "C" {
     PyObject *prefix##_qualname;                                            \
     _PyErr_StackItem prefix##_exc_state;                                    \
     PyObject *prefix##_origin_or_finalizer;                                 \
+    _PyContextChain _ctx_chain;                                             \
     char prefix##_hooks_inited;                                             \
     char prefix##_closed;                                                   \
     char prefix##_running_async;                                            \

@@ -9,6 +9,7 @@ extern "C" {
 #endif
 
 #include "pycore_brc.h"             // struct _brc_thread_state
+#include "pycore_contextchain.h"    // _PyContextChain
 #include "pycore_freelist_state.h"  // struct _Py_freelists
 #include "pycore_mimalloc.h"        // struct _mimalloc_thread_state
 #include "pycore_qsbr.h"            // struct qsbr
@@ -21,6 +22,9 @@ typedef struct _PyThreadStateImpl {
     // semi-public fields are in PyThreadState.
     PyThreadState base;
 
+    // Lazily initialized (must be zeroed at startup).
+    _PyContextChain _ctx_chain;
+
     PyObject *asyncio_running_loop; // Strong reference
 
     struct _qsbr_thread_state *qsbr;  // only used by free-threaded build