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

Author: rhansen

Doc/library/contextvars.rst

@@ -144,16 +144,72 @@ Manual Context Management
    To get a copy of the current context use the
    :func:`~contextvars.copy_context` function.
 
+   When *entering* a Context, either by calling the :meth:`Context.run` method
+   or using the Context object as a :term:`context manager`, the Context becomes
+   the *current Context*.  When *exiting* the current Context, either by
+   returning from the callback passed to :meth:`Context.run` or by exiting the
+   :keyword:`with` statement suite, the current Context reverts back to what it
+   was before the exited Context was entered.
+
+   Attempting to do any of the following will raise a :exc:`RuntimeError`:
+
+   * Entering an already entered Context.
+   * Exiting from a Context that is not the current Context.
+   * Exiting a Context from a different thread than the one used to enter the
+     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.  Thus, exiting
+   a Context effectively reverts any changes made to context variables while the
+   Context was entered.  (If desired, the values can be restored by re-entering
+   the Context.)
+
    Context implements the :class:`collections.abc.Mapping` interface.
 
+   .. versionadded:: 3.12
+      A Context object can be used as a :term:`context manager`.  The
+      :meth:`Context.__enter__` and :meth:`Context.__exit__` methods
+      (automatically called by the :keyword:`with` statement) enters and exits
+      the Context, respectively.  The value returned from
+      :meth:`Context.__enter__`, and thus bound to the identifier given in the
+      :keyword:`with` statement's :keyword:`!as` clause if present, is the
+      Context object itself.
+
+   Example:
+
+   .. testcode::
+
+      import contextvars
+
+      var = contextvars.ContextVar("var")
+      var.set("initial")
+
+      # Copy the current Context and enter it.
+      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 value of var should have reverted.
+      assert var.get() == "initial"
+      # But the updated value is still recorded in ctx.
+      assert ctx[var] == "updated"
+
+      # Re-entering ctx should restore 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')
@@ -181,10 +237,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.

Lib/test/test_context.py

@@ -3,6 +3,7 @@
 import functools
 import gc
 import random
+import threading
 import time
 import unittest
 import weakref
@@ -360,6 +361,62 @@ def sub(num):
             tp.shutdown()
         self.assertEqual(results, list(range(10)))
 
+    @isolated_context
+    def test_context_manager_1(self):
+        cvar = contextvars.ContextVar('cvar', default='initial')
+        self.assertEqual(cvar.get(), 'initial')
+        ctx = contextvars.copy_context()
+        with ctx as tmp:
+            self.assertIs(tmp, ctx)
+            self.assertEqual(cvar.get(), 'initial')
+            cvar.set('updated')
+            self.assertEqual(cvar.get(), 'updated')
+        self.assertEqual(cvar.get(), 'initial')
+        self.assertEqual(ctx[cvar], 'updated')
+        with ctx:
+            self.assertEqual(cvar.get(), 'updated')
+
+    @threading_helper.requires_working_threading()
+    def test_context_manager_2(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_3(self):
+        with contextvars.copy_context() as ctx:
+            with self.assertRaises(RuntimeError):
+                with ctx:
+                    pass
+
+    def test_context_manager_4(self):
+        with contextvars.copy_context() as ctx:
+            with self.assertRaises(RuntimeError):
+                ctx.run(lambda: None)
+
+    def test_context_manager_5(self):
+        ctx = contextvars.copy_context()
+
+        def fn():
+            with self.assertRaises(RuntimeError):
+                with ctx:
+                    pass
+
+        ctx.run(fn)
+
+    def test_context_manager_6(self):
+        with contextvars.copy_context() as ctx:
+            with contextvars.copy_context():
+                with self.assertRaises(RuntimeError):
+                    ctx.__exit__(None, None, None)
+
+    def test_context_manager_7(self):
+        ctx = contextvars.copy_context()
+        with self.assertRaises(RuntimeError):
+            ctx.__exit__(None, None, None)
+
 
 # HAMT Tests
 

Misc/ACKS

@@ -691,6 +691,7 @@ Manus Hand
 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 @@
+Add context manager methods to :class:`contextvars.Context`.

Python/clinic/context.c.h

@@ -153,7 +153,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");
@@ -558,6 +557,47 @@ context_tp_contains(PyContext *self, PyObject *key)
 }
 
 
+/*[clinic input]
+_contextvars.Context.__enter__
+
+Context manager enter.
+[clinic start generated code]*/
+
+static PyObject *
+_contextvars_Context___enter___impl(PyContext *self)
+/*[clinic end generated code: output=7374aea8983b777a input=093b5c4808ddf1b1]*/
+{
+    PyThreadState *ts = _PyThreadState_GET();
+    if (_PyContext_Enter(ts, (PyObject *)self)) {
+        return NULL;
+    }
+    return Py_NewRef(self);
+}
+
+
+/*[clinic input]
+_contextvars.Context.__exit__
+    exc_type: object
+    exc_val: object
+    exc_tb: object
+    /
+
+Context manager exit.
+[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=31b822b221c4439b]*/
+{
+    PyThreadState *ts = _PyThreadState_GET();
+    if (_PyContext_Exit(ts, (PyObject *)self)) {
+        return NULL;
+    }
+    Py_RETURN_FALSE;
+}
+
+
 /*[clinic input]
 _contextvars.Context.get
     key: object
@@ -677,6 +717,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