From aaceed08d8b15800be97946c41355bd12f86d62f Mon Sep 17 00:00:00 2001 From: Jason Grout Date: Tue, 11 Aug 2026 16:33:01 -0700 Subject: [PATCH 1/6] Enable temporarily setting a parent header that is only thread-local. --- ipykernel/displayhook.py | 41 ++++++++++++++++++++++++------ ipykernel/iostream.py | 12 +++++++-- ipykernel/zmqshell.py | 55 ++++++++++++++++++++++++++++++++-------- 3 files changed, 88 insertions(+), 20 deletions(-) diff --git a/ipykernel/displayhook.py b/ipykernel/displayhook.py index 9ec44747d..2d92fb170 100644 --- a/ipykernel/displayhook.py +++ b/ipykernel/displayhook.py @@ -64,11 +64,23 @@ 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 +100,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 +136,23 @@ 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.""" diff --git a/ipykernel/iostream.py b/ipykernel/iostream.py index a5b5b6c0b..6a66f6ef9 100644 --- a/ipykernel/iostream.py +++ b/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): diff --git a/ipykernel/zmqshell.py b/ipykernel/zmqshell.py index 67395e4e9..bd5c31c3f 100644 --- a/ipykernel/zmqshell.py +++ b/ipykernel/zmqshell.py @@ -20,6 +20,7 @@ import threading import typing import warnings +from contextlib import contextmanager from pathlib import Path from subprocess import CalledProcessError @@ -83,11 +84,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""" @@ -725,15 +737,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 +761,33 @@ 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] + if hasattr(self, "_data_pub"): + objs.append(self.data_pub) + objs += [sys.stdout, sys.stderr] + for obj in objs: + set_thread = getattr(obj, "set_thread_parent", None) + if set_thread is not None: + tokens.append((obj.reset_thread_parent, 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) + + @contextmanager + def parent_override(self, parent): + """Context manager to temporarily set the thread's parent header""" + tokens = self.set_thread_parent(parent) + try: + yield + finally: + self.reset_thread_parent(tokens) + def init_magics(self): """Initialize magics.""" super().init_magics() From 0c10601614f133ed60f4b1caa9223d60b24e1aad Mon Sep 17 00:00:00 2001 From: Jason Grout Date: Wed, 12 Aug 2026 04:53:33 -0700 Subject: [PATCH 2/6] Lint --- ipykernel/displayhook.py | 2 -- ipykernel/zmqshell.py | 2 +- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/ipykernel/displayhook.py b/ipykernel/displayhook.py index 2d92fb170..3b80e9a80 100644 --- a/ipykernel/displayhook.py +++ b/ipykernel/displayhook.py @@ -77,7 +77,6 @@ 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 global and thread parent header.""" self.parent_header = extract_header(parent) @@ -141,7 +140,6 @@ 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)) diff --git a/ipykernel/zmqshell.py b/ipykernel/zmqshell.py index bd5c31c3f..03e441897 100644 --- a/ipykernel/zmqshell.py +++ b/ipykernel/zmqshell.py @@ -764,7 +764,7 @@ def get_parent(self): 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] + objs: list[typing.Any] = [self.displayhook, self.display_pub] if hasattr(self, "_data_pub"): objs.append(self.data_pub) objs += [sys.stdout, sys.stderr] From 569cb925fc6f06d8c954bc63efedbe1891c3ec7a Mon Sep 17 00:00:00 2001 From: Jason Grout Date: Wed, 12 Aug 2026 13:55:57 -0700 Subject: [PATCH 3/6] Set a default _parent_header_global --- ipykernel/zmqshell.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/ipykernel/zmqshell.py b/ipykernel/zmqshell.py index 03e441897..69a3f3e23 100644 --- a/ipykernel/zmqshell.py +++ b/ipykernel/zmqshell.py @@ -547,14 +547,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): From 52247a29fc205c81e458b3f6ace3904e4796606c Mon Sep 17 00:00:00 2001 From: Jason Grout Date: Wed, 12 Aug 2026 15:04:49 -0700 Subject: [PATCH 4/6] Add set_thread_parent test --- tests/test_kernel.py | 30 +++++++++++++++++++++--------- 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/tests/test_kernel.py b/tests/test_kernel.py index db5debb28..2887c6fa6 100644 --- a/tests/test_kernel.py +++ b/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, + get_ipython().set_parent sets the thread-local value and global value, which supersedes the default. + get_ipython().set_thread_parent sets the thread-local value 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" From 4501fd9967f20b845c6ae072554a9e2c2bc2160b Mon Sep 17 00:00:00 2001 From: Jason Grout Date: Wed, 12 Aug 2026 16:12:51 -0700 Subject: [PATCH 5/6] Update tests --- tests/test_kernel.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_kernel.py b/tests/test_kernel.py index 2887c6fa6..fb9a7ad1f 100644 --- a/tests/test_kernel.py +++ b/tests/test_kernel.py @@ -125,10 +125,10 @@ def collect_outputs(get_iopub_msg, parent_msg_id, timeout=5): 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 and global 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 value for only the thread. + get_ipython().set_thread_parent sets the thread-local parent for only the thread. """ code = f"""\ from threading import Event, Thread @@ -232,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}) From cc87037ee42b174ed8d75d83964ccbfdddc083fc Mon Sep 17 00:00:00 2001 From: Jason Grout Date: Wed, 12 Aug 2026 23:12:29 -0700 Subject: [PATCH 6/6] Simplify and remove unused context manager --- ipykernel/zmqshell.py | 18 ++++-------------- 1 file changed, 4 insertions(+), 14 deletions(-) diff --git a/ipykernel/zmqshell.py b/ipykernel/zmqshell.py index 69a3f3e23..dd5216246 100644 --- a/ipykernel/zmqshell.py +++ b/ipykernel/zmqshell.py @@ -20,7 +20,6 @@ import threading import typing import warnings -from contextlib import contextmanager from pathlib import Path from subprocess import CalledProcessError @@ -766,14 +765,14 @@ def get_parent(self): 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: list[typing.Any] = [self.displayhook, self.display_pub] + objs = [self.displayhook, self.display_pub, sys.stdout, sys.stderr] if hasattr(self, "_data_pub"): objs.append(self.data_pub) - objs += [sys.stdout, sys.stderr] for obj in objs: set_thread = getattr(obj, "set_thread_parent", None) - if set_thread is not None: - tokens.append((obj.reset_thread_parent, set_thread(parent))) + 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): @@ -781,15 +780,6 @@ def reset_thread_parent(self, tokens): for reset, token in reversed(tokens): reset(token) - @contextmanager - def parent_override(self, parent): - """Context manager to temporarily set the thread's parent header""" - tokens = self.set_thread_parent(parent) - try: - yield - finally: - self.reset_thread_parent(tokens) - def init_magics(self): """Initialize magics.""" super().init_magics()