fix ContextVar persistence across cells (#1462)

Author: minrk

ipykernel/kernelbase.py

@@ -60,7 +60,7 @@
 
 from ._version import kernel_protocol_version
 from .iostream import OutStream
-from .utils import LazyDict
+from .utils import LazyDict, _async_in_context
 
 _AWAITABLE_MESSAGE: str = (
     "For consistency across implementations, it is recommended that `{func_name}`"
@@ -557,7 +557,7 @@ def start(self):
                 self.shell_stream.on_recv(self.shell_channel_thread_main, copy=False)
             else:
                 self.shell_stream.on_recv(
-                    partial(self.shell_main, None),
+                    _async_in_context(partial(self.shell_main, None)),
                     copy=False,
                 )
 

ipykernel/subshell_manager.py

@@ -15,6 +15,7 @@
 from .socket_pair import SocketPair
 from .subshell import SubshellThread
 from .thread import SHELL_CHANNEL_THREAD_NAME
+from .utils import _async_in_context
 
 
 class SubshellManager:
@@ -129,7 +130,9 @@ def set_on_recv_callback(self, on_recv_callback):
         """
         assert current_thread() == self._parent_thread
         self._on_recv_callback = on_recv_callback
-        self._shell_channel_to_main.on_recv(IOLoop.current(), partial(self._on_recv_callback, None))
+        self._shell_channel_to_main.on_recv(
+            IOLoop.current(), _async_in_context(partial(on_recv_callback, None))
+        )
 
     def set_subshell_aborting(self, subshell_id: str, aborting: bool) -> None:
         """Set the aborting flag of the specified subshell."""
@@ -165,7 +168,7 @@ def _create_subshell(self) -> str:
 
         subshell_thread.shell_channel_to_subshell.on_recv(
             subshell_thread.io_loop,
-            partial(self._on_recv_callback, subshell_id),
+            _async_in_context(partial(self._on_recv_callback, subshell_id)),
         )
 
         subshell_thread.subshell_to_shell_channel.on_recv(

ipykernel/utils.py

@@ -1,7 +1,17 @@
 """Utilities"""
 
+from __future__ import annotations
+
+import asyncio
+import sys
 import typing as t
 from collections.abc import Mapping
+from contextvars import copy_context
+from functools import partial, wraps
+
+if t.TYPE_CHECKING:
+    from collections.abc import Callable
+    from contextvars import Context
 
 
 class LazyDict(Mapping[str, t.Any]):
@@ -24,3 +34,48 @@ def __len__(self):
 
     def __iter__(self):
         return iter(self._dict)
+
+
+T = t.TypeVar("T")
+U = t.TypeVar("U")
+V = t.TypeVar("V")
+
+
+def _async_in_context(
+    f: Callable[..., t.Coroutine[T, U, V]], context: Context | None = None
+) -> Callable[..., t.Coroutine[T, U, V]]:
+    """
+    Wrapper to run a coroutine in a persistent ContextVar Context.
+
+    Backports asyncio.create_task(context=...) behavior from Python 3.11
+    """
+    if context is None:
+        context = copy_context()
+
+    if sys.version_info >= (3, 11):
+
+        @wraps(f)
+        async def run_in_context(*args, **kwargs):
+            coro = f(*args, **kwargs)
+            return await asyncio.create_task(coro, context=context)
+
+        return run_in_context
+
+    # don't need this backport when we require 3.11
+    # context_holder so we have a modifiable container for later calls
+    context_holder = [context]  # type: ignore[unreachable]
+
+    async def preserve_context(f, *args, **kwargs):
+        """call a coroutine, preserving the context after it is called"""
+        try:
+            return await f(*args, **kwargs)
+        finally:
+            # persist changes to the context for future calls
+            context_holder[0] = copy_context()
+
+    @wraps(f)
+    async def run_in_context_pre311(*args, **kwargs):
+        ctx = context_holder[0]
+        return await ctx.run(partial(asyncio.create_task, preserve_context(f, *args, **kwargs)))
+
+    return run_in_context_pre311

ipykernel/zmqshell.py

@@ -570,7 +570,7 @@ def _update_exit_now(self, change):
     # Over ZeroMQ, GUI control isn't done with PyOS_InputHook as there is no
     # interactive input being read; we provide event loop support in ipkernel
     def enable_gui(self, gui: typing.Any = None) -> None:
-        """Enable a given guil."""
+        """Enable a given gui."""
         from .eventloops import enable_gui as real_enable_gui
 
         try:

tests/test_kernel.py

@@ -848,3 +848,33 @@ def test_parent_header_and_ident():
         msg_id, _ = execute(kc=kc, code="print(k._parent_ident['control'])")
         stdout, _ = assemble_output(kc.get_iopub_msg, parent_msg_id=msg_id)
         assert stdout == f"[b'{session}']\n"
+
+
+def test_context_vars():
+    with new_kernel() as kc:
+        msg_id, _ = execute(
+            kc=kc,
+            code="from contextvars import ContextVar, copy_context\nctxvar = ContextVar('var', default='default')",
+        )
+        stdout, _ = assemble_output(kc.get_iopub_msg, parent_msg_id=msg_id)
+
+        msg_id, _ = execute(
+            kc=kc,
+            code="print(ctxvar.get())",
+        )
+        stdout, _ = assemble_output(kc.get_iopub_msg, parent_msg_id=msg_id)
+        assert stdout.strip() == "default"
+
+        msg_id, _ = execute(
+            kc=kc,
+            code="ctxvar.set('set'); print(ctxvar.get())",
+        )
+        stdout, _ = assemble_output(kc.get_iopub_msg, parent_msg_id=msg_id)
+        assert stdout.strip() == "set"
+
+        msg_id, _ = execute(
+            kc=kc,
+            code="print(ctxvar.get())",
+        )
+        stdout, _ = assemble_output(kc.get_iopub_msg, parent_msg_id=msg_id)
+        assert stdout.strip() == "set"