diff --git a/ipykernel/displayhook.py b/ipykernel/displayhook.py index 9ec44747d..3b80e9a80 100644 --- a/ipykernel/displayhook.py +++ b/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.""" 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..dd5216246 100644 --- a/ipykernel/zmqshell.py +++ b/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() diff --git a/tests/test_kernel.py b/tests/test_kernel.py index db5debb28..fb9a7ad1f 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, - 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})