pythongh-99633: Add context manager support to contextvars.Context

Doc/library/contextvars.rst

@@ -144,21 +144,79 @@ Manual Context Management
    To get a copy of the current context use the
    :func:`~contextvars.copy_context` function.
 
-   Every thread will have a different top-level :class:`~contextvars.Context`
-   object. This means that a :class:`ContextVar` object behaves in a similar
-   fashion to :func:`threading.local` when values are assigned in different
-   threads.
+   Each thread has its own effective stack of :class:`!Context` objects.  The
+   *current context* is the :class:`!Context` object at the top of the current
+   thread's stack.  All :class:`!Context` objects in the stacks are considered
+   to be *entered*.
+
+   *Entering* a context, either by calling the :meth:`~Context.run` method or
+   using the context as a :term:`context manager`, pushes the context onto the
+   top of the current thread's stack, making it the current context.
+
+   *Exiting* from the current context, either by returning from the callback
+   passed to :meth:`~Context.run` or by exiting the :keyword:`with` statement
+   suite, pops the context off of the top of the stack, restoring the current
+   context to what it was before.
+
+   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 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).
+
+   Any changes to :class:`ContextVar` values via the :meth:`ContextVar.set`
+   method are recorded in the current context.  The :meth:`ContextVar.get`
+   method returns the value associated with the current context.  Exiting a
+   context effectively reverts any changes made to context variables while the
+   context was entered (if needed, the values can be restored by re-entering the
+   context).
 
    Context implements the :class:`collections.abc.Mapping` interface.
 
+   .. versionadded:: 3.14
+      Added support for the :term:`context management protocol`.  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")
+         assert var.get() == "initial"
+
+         # Copy the current Context and enter the copy.
+         with contextvars.copy_context() as ctx:
+             var.set("updated")
+             assert var in ctx
+             assert ctx[var] == "updated"
+             assert var.get() == "updated"
+
+         # Exited ctx, so the observed value of var has reverted.
+         assert var.get() == "initial"
+         # But the updated value is still recorded in ctx.
+         assert ctx[var] == "updated"
+
+         # Re-entering ctx restores the updated value of var.
+         with ctx:
+             assert var.get() == "updated"
+
    .. method:: run(callable, *args, **kwargs)
 
-      Execute ``callable(*args, **kwargs)`` code in the context object
-      the *run* method is called on.  Return the result of the execution
-      or propagate an exception if one occurred.
+      Enters the Context, executes ``callable(*args, **kwargs)``, then exits the
+      Context.  Returns *callable*'s return value, or propagates an exception if
+      one occurred.
 
-      Any changes to any context variables that *callable* makes will
-      be contained in the context object::
+      Example::
 
         var = ContextVar('var')
         var.set('spam')
@@ -186,10 +244,6 @@ Manual Context Management
         # However, outside of 'ctx', 'var' is still set to 'spam':
         # var.get() == 'spam'
 
-      The method raises a :exc:`RuntimeError` when called on the same
-      context object from more than one OS thread, or when called
-      recursively.
-
    .. method:: copy()
 
       Return a shallow copy of the context object.

Doc/whatsnew/3.14.rst

@@ -112,6 +112,13 @@ ast
   (Contributed by Bénédikt Tran in :gh:`121141`.)
 
 
+contextvars
+-----------
+
+* Added support for the :term:`context management protocol` to
+  :class:`contextvars.Context`.  (Contributed by Richard Hansen in :gh:`99634`.)
+
+
 ctypes
 ------
 

Lib/test/test_context.py

@@ -3,6 +3,7 @@
 import functools
 import gc
 import random
+import threading
 import time
 import unittest
 import weakref
@@ -361,6 +362,79 @@ def sub(num):
             tp.shutdown()
         self.assertEqual(results, list(range(10)))
 
+    @isolated_context
+    def test_context_manager(self):
+        cvar = contextvars.ContextVar('cvar', default='initial')
+        self.assertEqual(cvar.get(), 'initial')
+        with contextvars.copy_context():
+            self.assertEqual(cvar.get(), 'initial')
+            cvar.set('updated')
+            self.assertEqual(cvar.get(), 'updated')
+        self.assertEqual(cvar.get(), 'initial')
+
+    def test_context_manager_as_binding(self):
+        ctx = contextvars.copy_context()
+        with ctx as ctx_as_binding:
+            self.assertIs(ctx_as_binding, ctx)
+
+    def test_context_manager_nested(self):
+        with contextvars.copy_context() as outer_ctx:
+            with contextvars.copy_context() as inner_ctx:
+                self.assertIsNot(outer_ctx, inner_ctx)
+
+    @isolated_context
+    def test_context_manager_enter_again_after_exit(self):
+        cvar = contextvars.ContextVar('cvar', default='initial')
+        self.assertEqual(cvar.get(), 'initial')
+        with contextvars.copy_context() as ctx:
+            cvar.set('updated')
+            self.assertEqual(cvar.get(), 'updated')
+        self.assertEqual(cvar.get(), 'initial')
+        with ctx:
+            self.assertEqual(cvar.get(), 'updated')
+        self.assertEqual(cvar.get(), 'initial')
+
+    @threading_helper.requires_working_threading()
+    def test_context_manager_rejects_exit_from_different_thread(self):
+        ctx = contextvars.copy_context()
+        thread = threading.Thread(target=ctx.__enter__)
+        thread.start()
+        thread.join()
+        with self.assertRaises(RuntimeError):
+            ctx.__exit__(None, None, None)
+
+    def test_context_manager_is_not_reentrant(self):
+        with self.subTest('context manager then context manager'):
+            with contextvars.copy_context() as ctx:
+                with self.assertRaises(RuntimeError):
+                    with ctx:
+                        pass
+        with self.subTest('context manager then run method'):
+            with contextvars.copy_context() as ctx:
+                with self.assertRaises(RuntimeError):
+                    ctx.run(lambda: None)
+        with self.subTest('run method then context manager'):
+            ctx = contextvars.copy_context()
+
+            def fn():
+                with self.assertRaises(RuntimeError):
+                    with ctx:
+                        pass
+
+            ctx.run(fn)
+
+    def test_context_manager_rejects_noncurrent_exit(self):
+        with contextvars.copy_context() as outer_ctx:
+            with contextvars.copy_context() as inner_ctx:
+                self.assertIsNot(outer_ctx, inner_ctx)
+                with self.assertRaises(RuntimeError):
+                    outer_ctx.__exit__(None, None, None)
+
+    def test_context_manager_rejects_nonentered_exit(self):
+        ctx = contextvars.copy_context()
+        with self.assertRaises(RuntimeError):
+            ctx.__exit__(None, None, None)
+
 
 # HAMT Tests
 

Misc/ACKS

@@ -716,6 +716,7 @@ Michael Handler
 Andreas Hangauer
 Milton L. Hankins
 Carl Bordum Hansen
+Richard Hansen
 Stephen Hansen
 Barry Hantman
 Lynda Hardman

Misc/NEWS.d/next/Library/2022-11-21-01-24-46.gh-issue-99633.vhrNRe.rst

@@ -0,0 +1,2 @@
+Added support for the :term:`context management protocol` to
+:class:`contextvars.Context`.  Patch by Richard Hansen.

Python/context.c

@@ -144,7 +144,6 @@ _PyContext_Exit(PyThreadState *ts, PyObject *octx)
     }
 
     if (ts->context != (PyObject *)ctx) {
-        /* Can only happen if someone misuses the C API */
         PyErr_SetString(PyExc_RuntimeError,
                         "cannot exit context: thread state references "
                         "a different context object");
@@ -521,6 +520,67 @@ context_tp_contains(PyContext *self, PyObject *key)
 }
 
 
+/*[clinic input]
+_contextvars.Context.__enter__
+
+Context manager enter.
+
+Automatically called by the 'with' statement.  Using the Context object as a
+context manager is an alternative to calling the Context.run() method.
+
+Example:
+
+    var = contextvars.ContextVar('var')
+    var.set('initial')
+
+    with contextvars.copy_context():
+        # Changes to context variables will be rolled back upon exiting the
+        # `with` statement.
+        var.set('updated')
+        print(var.get())  # 'updated'
+
+    # The context variable value has been rolled back.
+    print(var.get())  # 'initial'
+[clinic start generated code]*/
+
+static PyObject *
+_contextvars_Context___enter___impl(PyContext *self)
+/*[clinic end generated code: output=7374aea8983b777a input=fffe71e56ca17ee4]*/
+{
+    // The new ref added here is for the `with` statement's `as` binding.  It is
+    // decremented when the variable goes out of scope, which can be before or
+    // after `PyContext_Exit` is called.  (The binding can go out of scope
+    // immediately -- before the `with` suite even runs -- if there is no `as`
+    // clause.  Or it can go out of scope long after the `with` suite completes
+    // because `with` does not have its own scope.)  Because of this timing, two
+    // references are needed: the one added in `PyContext_Enter` and the one
+    // added here.
+    return PyContext_Enter((PyObject *)self) < 0 ? NULL : PyNewRef(self);
+}
+
+
+/*[clinic input]
+_contextvars.Context.__exit__
+    exc_type: object
+    exc_val: object
+    exc_tb: object
+    /
+
+Context manager exit.
+
+Automatically called at the conclusion of a 'with' statement when the Context is
+used as a context manager.  See the Context.__enter__() method for more details.
+[clinic start generated code]*/
+
+static PyObject *
+_contextvars_Context___exit___impl(PyContext *self, PyObject *exc_type,
+                                   PyObject *exc_val, PyObject *exc_tb)
+/*[clinic end generated code: output=4608fa9151f968f1 input=ff70cbbf6a112b1d]*/
+{
+    return PyContext_Exit((PyObject *)self) < 0 ? NULL : Py_None;
+}
+
+
 /*[clinic input]
 _contextvars.Context.get
     key: object
@@ -641,6 +701,8 @@ context_run(PyContext *self, PyObject *const *args,
 
 
 static PyMethodDef PyContext_methods[] = {
+    _CONTEXTVARS_CONTEXT___ENTER___METHODDEF
+    _CONTEXTVARS_CONTEXT___EXIT___METHODDEF
     _CONTEXTVARS_CONTEXT_GET_METHODDEF
     _CONTEXTVARS_CONTEXT_ITEMS_METHODDEF
     _CONTEXTVARS_CONTEXT_KEYS_METHODDEF