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

rhansen · rhansen · commit 01afd3c227de · 2024-10-10T04:42:33.000-04:00
diff --git a/Doc/library/contextvars.rst b/Doc/library/contextvars.rst
@@ -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
 ------
 
diff --git a/Lib/test/test_context.py b/Lib/test/test_context.py
@@ -1,8 +1,11 @@
+import asyncio
 import concurrent.futures
+import contextlib
 import contextvars
 import functools
 import gc
 import random
+import threading
 import time
 import unittest
 import weakref
@@ -369,6 +372,88 @@ 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)
+
+    @isolated_context
+    def test_context_manager_nested(self):
+        cvar = contextvars.ContextVar('cvar', default='default')
+        with contextvars.copy_context() as outer_ctx:
+            cvar.set('outer')
+            with contextvars.copy_context() as inner_ctx:
+                self.assertIsNot(outer_ctx, inner_ctx)
+                self.assertEqual(cvar.get(), 'outer')
+                cvar.set('inner')
+                self.assertEqual(outer_ctx[cvar], 'outer')
+                self.assertEqual(cvar.get(), 'inner')
+            self.assertEqual(cvar.get(), 'outer')
+        self.assertEqual(cvar.get(), 'default')
+
+    @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)
+
 
 class GeneratorContextTest(unittest.TestCase):
     def test_default_is_none(self):
@@ -641,6 +726,145 @@ def makegen_outer():
             self.assertEqual(gen_outer.send(ctx_outer), ('inner', 'outer'))
             self.assertEqual(cvar.get(), 'updated')
 
+
+class AsyncContextTest(unittest.IsolatedAsyncioTestCase):
+    async def test_asyncio_independent_contexts(self):
+        """Check that coroutines are run with independent contexts.
+
+        Changes to context variables outside a coroutine should not affect the
+        values seen inside the coroutine and vice-versa.  (This might be
+        implemented by manually setting the context before executing each step
+        of (send to) a coroutine, or by ensuring that the coroutine is an
+        independent coroutine before executing any steps.)
+        """
+        cvar = contextvars.ContextVar('cvar', default='A')
+        updated1 = asyncio.Event()
+        updated2 = asyncio.Event()
+
+        async def task1():
+            self.assertIs(cvar.get(), 'A')
+            await asyncio.sleep(0)
+            cvar.set('B')
+            await asyncio.sleep(0)
+            updated1.set()
+            await updated2.wait()
+            self.assertIs(cvar.get(), 'B')
+
+        async def task2():
+            await updated1.wait()
+            self.assertIs(cvar.get(), 'A')
+            await asyncio.sleep(0)
+            cvar.set('C')
+            await asyncio.sleep(0)
+            updated2.set()
+            await asyncio.sleep(0)
+            self.assertIs(cvar.get(), 'C')
+
+        async with asyncio.TaskGroup() as tg:
+            tg.create_task(task1())
+            tg.create_task(task2())
+
+        self.assertIs(cvar.get(), 'A')
+
+    async def test_asynccontextmanager_is_dependent_by_default(self):
+        """Async generator in asynccontextmanager is dependent by default.
+
+        Context switches during the yield of a generator wrapped with
+        contextlib.asynccontextmanager should be visible to the generator by
+        default (for backwards compatibility).
+        """
+        cvar = contextvars.ContextVar('cvar', default='A')
+
+        @contextlib.asynccontextmanager
+        async def makecm():
+            await asyncio.sleep(0)
+            self.assertEqual(cvar.get(), 'A')
+            await asyncio.sleep(0)
+            # Everything above runs during __aenter__.
+            yield cvar.get()
+            # Everything below runs during __aexit__.
+            await asyncio.sleep(0)
+            self.assertEqual(cvar.get(), 'C')
+            await asyncio.sleep(0)
+            cvar.set('D')
+            await asyncio.sleep(0)
+
+        cm = makecm()
+        val = await cm.__aenter__()
+        self.assertEqual(val, 'A')
+        self.assertEqual(cvar.get(), 'A')
+        cvar.set('B')
+
+        with contextvars.copy_context():
+            cvar.set('C')
+            await cm.__aexit__(None, None, None)
+            self.assertEqual(cvar.get(), 'D')
+        self.assertEqual(cvar.get(), 'B')
+
+    async def test_asynccontextmanager_independent(self):
+        cvar = contextvars.ContextVar('cvar', default='A')
+
+        @contextlib.asynccontextmanager
+        async def makecm():
+            # Context.__enter__ called from a generator makes the generator
+            # independent while the `with` statement suite runs.
+            # (Alternatively we could have set the generator's _context
+            # property.)
+            with contextvars.copy_context():
+                await asyncio.sleep(0)
+                self.assertEqual(cvar.get(), 'A')
+                await asyncio.sleep(0)
+                # Everything above runs during __aenter__.
+                yield cvar.get()
+                # Everything below runs during __aexit__.
+                await asyncio.sleep(0)
+                self.assertEqual(cvar.get(), 'A')
+                await asyncio.sleep(0)
+                cvar.set('D')
+                await asyncio.sleep(0)
+
+        cm = makecm()
+        val = await cm.__aenter__()
+        self.assertEqual(val, 'A')
+        self.assertEqual(cvar.get(), 'A')
+        cvar.set('B')
+        with contextvars.copy_context():
+            cvar.set('C')
+            await cm.__aexit__(None, None, None)
+            self.assertEqual(cvar.get(), 'C')
+        self.assertEqual(cvar.get(), 'B')
+
+    async def test_generator_switch_between_independent_dependent(self):
+        cvar = contextvars.ContextVar('cvar', default='default')
+        with contextvars.copy_context() as ctx1:
+            cvar.set('in ctx1')
+        with contextvars.copy_context() as ctx2:
+            cvar.set('in ctx2')
+        with contextvars.copy_context() as ctx3:
+            cvar.set('in ctx3')
+
+        async def makegen():
+            await asyncio.sleep(0)
+            yield cvar.get()
+            await asyncio.sleep(0)
+            yield cvar.get()
+            await asyncio.sleep(0)
+            with ctx2:
+                yield cvar.get()
+                await asyncio.sleep(0)
+                yield cvar.get()
+                await asyncio.sleep(0)
+            yield cvar.get()
+
+        gen = makegen()
+        self.assertEqual(await anext(gen), 'default')
+        with ctx1:
+            self.assertEqual(await anext(gen), 'in ctx1')
+            self.assertEqual(await anext(gen), 'in ctx2')
+            with ctx3:
+                self.assertEqual(await anext(gen), 'in ctx2')
+            self.assertEqual(await anext(gen), 'in ctx1')
+
 # HAMT Tests
 
 
diff --git a/Misc/ACKS b/Misc/ACKS
@@ -716,6 +716,7 @@ Michael Handler
 Andreas Hangauer
 Milton L. Hankins
 Carl Bordum Hansen
+Richard Hansen
 Stephen Hansen
 Barry Hantman
 Lynda Hardman
diff --git a/Misc/NEWS.d/next/Library/2022-11-21-01-24-46.gh-issue-99633.vhrNRe.rst b/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.
diff --git a/Python/clinic/context.c.h b/Python/clinic/context.c.h
diff --git a/Python/context.c b/Python/context.c

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