Add API for changing the context of a running task.

docs/source/reference-lowlevel.rst

@@ -540,6 +540,20 @@ Task API
       :func:`wait_task_rescheduled` for details.)
 
 
+Low-level context tools
+=======================
+
+You can change the current :class:`contextvars.Context` a task is running
+in with a helper context manager.
+
+.. autofunction:: set_current_context
+    :async-with:
+
+.. note::
+
+    This context manager introduces a scheduling point but not a cancellation
+    point.
+
 .. _guest-mode:
 
 Using "guest mode" to run Trio on top of other event loops

newsfragments/2086.feature.rst

@@ -0,0 +1,2 @@
+Add the :func:`trio.lowlevel.set_current_context` API to allow changing the
+current :class:`contextvars.Context` a task is running in.

trio/_core/__init__.py

@@ -75,6 +75,8 @@
 
 from ._mock_clock import MockClock
 
+from ._context import set_current_context
+
 # Windows imports
 if sys.platform == "win32":
     from ._run import (

trio/_core/_context.py

@@ -0,0 +1,21 @@
+from async_generator import asynccontextmanager
+
+from . import _run
+
+
+@asynccontextmanager
+async def set_current_context(context):
+    """Returns an asynchronous context manager to change the
+    :class:`contextvars.Context` for the current task.
+
+    """
+    task = _run.current_task()
+    saved_context = task.context
+    task.context = context
+
+    try:
+        await _run.cancel_shielded_checkpoint()
+        yield
+    finally:
+        task.context = saved_context
+        await _run.cancel_shielded_checkpoint()

trio/_core/tests/test_context.py

@@ -0,0 +1,31 @@
+from contextvars import ContextVar, copy_context
+
+from .._run import CancelScope
+from .._context import set_current_context
+
+var = ContextVar("var")
+
+
+async def test_context_change():
+    var.set("abc")
+
+    new_context = copy_context()
+    async with set_current_context(new_context):
+        assert var.get() == "abc"
+        var.set("def")
+        assert var.get() == "def"
+
+    assert var.get() == "abc"
+
+
+async def test_context_change_with_cancellation():
+    var.set("qqq")
+
+    new_context = copy_context()
+    with CancelScope() as scope:
+        async with set_current_context(new_context):
+            var.set("ghj")
+            assert var.get() == "ghj"
+            scope.cancel()
+
+    assert var.get() == "qqq"

trio/lowlevel.py

@@ -44,6 +44,7 @@
     notify_closing,
     start_thread_soon,
     start_guest_run,
+    set_current_context,
 )
 
 from ._subprocess import open_process