Address review comments.

docs/source/reference-core.rst

@@ -1020,12 +1020,6 @@ Example output (yours may differ slightly):
    request 0: Helper task b finished
    request 0: Request received finished
 
-You can change the current :class:`contextvars.Context` a task is running
-in with a helper context manager.
-
-.. autofunction:: change_context
-    :async-with:
-
 For more information, read the
 `contextvar docs <https://docs.python.org/3.7/library/contextvars.html>`__.
 

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

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

trio/__init__.py

@@ -32,7 +32,6 @@
     BrokenResourceError,
     EndOfChannel,
     Nursery,
-    change_context,
 )
 
 from ._timeouts import (

trio/_core/__init__.py

@@ -75,7 +75,7 @@
 
 from ._mock_clock import MockClock
 
-from ._context import change_context
+from ._context import set_current_context
 
 # Windows imports
 if sys.platform == "win32":

trio/_core/_context.py

@@ -1,24 +1,21 @@
-from contextvars import Context
-
 from async_generator import asynccontextmanager
 
 from . import _run
 
 
 @asynccontextmanager
-async def change_context(context):
-    """Asynchronous context manager to change the :class:`contextvars.Context`
-    for the current task.
+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.checkpoint()
+        await _run.cancel_shielded_checkpoint()
         yield
     finally:
         task.context = saved_context
-        # I assume this is right?
         await _run.cancel_shielded_checkpoint()

trio/_core/tests/test_context.py

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

trio/lowlevel.py

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