Enable setting and resetting a thread-local parent header

ipykernel/displayhook.py

@@ -64,11 +64,22 @@ def parent_header(self):
         except LookupError:
             return self._parent_header_global
 
+    @parent_header.setter
+    def parent_header(self, value):
+        self._parent_header.set(value)
+        self._parent_header_global = value
+
+    def set_thread_parent(self, parent):
+        """Set the parent header for the calling thread only. Returns a reset token that can be used with reset_thread_parent."""
+        return self._parent_header.set(extract_header(parent))
+
+    def reset_thread_parent(self, token):
+        """Reset the parent header to undo the set_thread_parent call that returned the token."""
+        self._parent_header.reset(token)
+
     def set_parent(self, parent):
-        """Set the parent header."""
-        parent_header = extract_header(parent)
-        self._parent_header.set(parent_header)
-        self._parent_header_global = parent_header
+        """Set the global and thread parent header."""
+        self.parent_header = extract_header(parent)
 
 
 class ZMQShellDisplayHook(DisplayHook):
@@ -88,6 +99,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._parent_header = ContextVar("parent_header")
         self._parent_header.set({})
+        self._parent_header_global = {}
 
     @default("_thread_local")
     def _default_thread_local(self):
@@ -123,11 +135,22 @@ def parent_header(self):
         except LookupError:
             return self._parent_header_global
 
+    @parent_header.setter
+    def parent_header(self, value):
+        self._parent_header.set(value)
+        self._parent_header_global = value
+
+    def set_thread_parent(self, parent):
+        """Set the parent header for the calling thread only. Returns a reset token that can be used with reset_thread_parent."""
+        return self._parent_header.set(extract_header(parent))
+
+    def reset_thread_parent(self, token):
+        """Reset the parent header to undo the set_thread_parent call that returned the token."""
+        self._parent_header.reset(token)
+
     def set_parent(self, parent):
-        """Set the parent header."""
-        parent_header = extract_header(parent)
-        self._parent_header.set(parent_header)
-        self._parent_header_global = parent_header
+        """Set the global and thread parent header."""
+        self.parent_header = extract_header(parent)
 
     def start_displayhook(self):
         """Start the display hook."""

ipykernel/iostream.py

@@ -605,8 +605,8 @@ def parent_header(self):
 
     @parent_header.setter
     def parent_header(self, value):
+        self._parent_header.set(value)
         self._parent_header_global = value
-        return self._parent_header.set(value)
 
     def isatty(self):
         """Return a bool indicating whether this is an 'interactive' stream.
@@ -632,8 +632,16 @@ def _setup_stream_redirects(self, name):
     def _is_master_process(self):
         return os.getpid() == self._master_pid
 
+    def set_thread_parent(self, parent):
+        """Set the parent header for the calling thread only. Returns a reset token that can be used with reset_thread_parent."""
+        return self._parent_header.set(extract_header(parent))
+
+    def reset_thread_parent(self, token):
+        """Reset the parent header to undo the set_thread_parent call that returned the token."""
+        self._parent_header.reset(token)
+
     def set_parent(self, parent):
-        """Set the parent header."""
+        """Set the global and thread parent header."""
         self.parent_header = extract_header(parent)
 
     def close(self):

ipykernel/zmqshell.py

@@ -83,11 +83,22 @@ def parent_header(self):
         except LookupError:
             return self._parent_header_global
 
+    @parent_header.setter
+    def parent_header(self, value):
+        self._parent_header.set(value)
+        self._parent_header_global = value
+
+    def set_thread_parent(self, parent):
+        """Set the parent header for the calling thread only. Returns a reset token that can be used with reset_thread_parent."""
+        return self._parent_header.set(extract_header(parent))
+
+    def reset_thread_parent(self, token):
+        """Reset the parent header to undo the set_thread_parent call that returned the token."""
+        self._parent_header.reset(token)
+
     def set_parent(self, parent):
-        """Set the parent for outbound messages."""
-        parent_header = extract_header(parent)
-        self._parent_header.set(parent_header)
-        self._parent_header_global = parent_header
+        """Set the global and thread parent header."""
+        self.parent_header = extract_header(parent)
 
     def _flush_streams(self):
         """flush IO Streams prior to display"""
@@ -535,14 +546,16 @@ def __init__(self, *args, **kwargs):
         if "IPKernelApp" not in self.config:
             self.config.IPKernelApp.tqdm = "dummy value for https://github.com/tqdm/tqdm/pull/1628"
 
-        self._parent_header = contextvars.ContextVar("parent_header")
+        self._parent_header: contextvars.ContextVar[dict[str, typing.Any]] = contextvars.ContextVar(
+            "parent_header"
+        )
         self._parent_header.set({})
+        self._parent_header_global = {}
 
     displayhook_class = Type(ZMQShellDisplayHook)
     display_pub_class = Type(ZMQDisplayPublisher)
     data_pub_class = Any()
     kernel = Any()
-    _parent_header: contextvars.ContextVar[dict[str, Any]]
 
     @default("banner1")
     def _default_banner1(self):
@@ -725,15 +738,11 @@ def parent_header(self):
 
     @parent_header.setter
     def parent_header(self, value):
-        self._parent_header_global = value
         self._parent_header.set(value)
+        self._parent_header_global = value
 
     def set_parent(self, parent):
-        """Set the parent header for associating output with its triggering input
-
-        When called from a thread, sets the thread-local value, which persists
-        until the next call from this thread.
-        """
+        """Set the global and thread parent header for associating output with its triggering input."""
         self.parent_header = parent
         self.displayhook.set_parent(parent)  # type:ignore[attr-defined]
         self.display_pub.set_parent(parent)  # type:ignore[attr-defined]
@@ -753,6 +762,24 @@ def get_parent(self):
         """
         return self.parent_header
 
+    def set_thread_parent(self, parent):
+        """Set the parent header for only the current thread associating output with its triggering input"""
+        tokens = [(self._parent_header.reset, self._parent_header.set(parent))]
+        objs = [self.displayhook, self.display_pub, sys.stdout, sys.stderr]
+        if hasattr(self, "_data_pub"):
+            objs.append(self.data_pub)
+        for obj in objs:
+            set_thread = getattr(obj, "set_thread_parent", None)
+            reset_thread = getattr(obj, "reset_thread_parent", None)
+            if set_thread is not None and reset_thread is not None:
+                tokens.append((reset_thread, set_thread(parent)))
+        return tuple(tokens)
+
+    def reset_thread_parent(self, tokens):
+        """Reset the parent header to undo the set_thread_parent call that returned the token."""
+        for reset, token in reversed(tokens):
+            reset(token)
+
     def init_magics(self):
         """Initialize magics."""
         super().init_magics()

tests/test_kernel.py

@@ -121,28 +121,31 @@ def collect_outputs(get_iopub_msg, parent_msg_id, timeout=5):
             print(msg["msg_type"])
 
 
-@pytest.mark.parametrize("explicit_parent", [True, False])
-def test_print_to_correct_cell_from_thread(explicit_parent: bool):
+@pytest.mark.parametrize("explicit_parent", ["global", "thread-local", "none"])
+def test_print_to_correct_cell_from_thread(explicit_parent: str):
     """should print to the current cell unless
 
-    get_ipython().set_parent sets the thread-local value,
-    which supersedes the default.
+    get_ipython().set_parent sets the thread-local parent and the global parent,
+    which supersedes the default parent set by the current shell execution.
 
+    get_ipython().set_thread_parent sets the thread-local parent for only the thread.
     """
     code = f"""\
         from threading import Event, Thread
         from time import sleep
         from IPython.display import display
 
-        explicit_parent = {explicit_parent}
+        explicit_parent = "{explicit_parent}"
         parent = get_ipython().get_parent()
 
         cell_start_event = Event()
         cell_end_event = Event()
 
         def thread_target():
-            if explicit_parent:
+            if explicit_parent == "global":
                 get_ipython().set_parent(parent)
+            elif explicit_parent == "thread-local":
+                reset_parent_token = get_ipython().set_thread_parent(parent)
 
             print("before", flush=True)
             display(1)
@@ -151,6 +154,8 @@ def thread_target():
 
             print("during", flush=True)
             display(2)
+            if explicit_parent == "thread-local":
+                get_ipython().reset_thread_parent(reset_parent_token)
             cell_end_event.set()
             cell_start_event.wait(timeout=10)
             cell_start_event.clear()
@@ -190,13 +195,20 @@ def add_output(msg):
         last_cell_msg_id = kc.execute("cell_start_event.set()\nthread.join()")
         for msg in collect_outputs(kc.get_iopub_msg, last_cell_msg_id):
             add_output(msg)
-    print(outputs)
-    if explicit_parent:
-        # assert next_cell_msg_id not in outputs
-        # assert last_cell_msg_id not in outputs
+    if explicit_parent == "global":
+        assert next_cell_msg_id not in outputs
+        assert last_cell_msg_id not in outputs
         thread_cell_output = outputs[thread_msg_id]
         assert thread_cell_output["stdout"] == "before\nduring\nafter\n"
         assert thread_cell_output["display_data"] == ["1", "2", "3"]
+    elif explicit_parent == "thread-local":
+        assert next_cell_msg_id not in outputs
+        thread_cell_output = outputs[thread_msg_id]
+        assert thread_cell_output["stdout"] == "before\nduring\n"
+        assert thread_cell_output["display_data"] == ["1", "2"]
+        last_cell_output = outputs[last_cell_msg_id]
+        assert last_cell_output["stdout"] == "after\n"
+        assert last_cell_output["display_data"] == ["3"]
     else:
         thread_cell_output = outputs[thread_msg_id]
         assert thread_cell_output["stdout"] == "before\n"
@@ -220,7 +232,7 @@ def test_print_to_correct_cell_from_child_thread():
     parent = get_ipython().get_parent()
 
     def child_target():
-        get_ipython().set_parent(parent)
+        get_ipython().set_thread_parent(parent)
         for i in range({iterations}):
             print(i, end='', flush=True)
             sleep({interval})