diff --git a/ipykernel/displayhook.py b/ipykernel/displayhook.py index 3b80e9a80..1c0e08b5d 100644 --- a/ipykernel/displayhook.py +++ b/ipykernel/displayhook.py @@ -28,8 +28,8 @@ def __init__(self, session, pub_socket): self.session = session self.pub_socket = pub_socket - self._parent_header: ContextVar[dict[str, Any]] = ContextVar("parent_header") - self._parent_header.set({}) + self._parent_header: ContextVar[tuple[int, dict[str, Any]]] = ContextVar("parent_header") + self._parent_header.set((threading.get_ident(), {})) self._parent_header_global = {} def get_execution_count(self): @@ -60,18 +60,21 @@ def __call__(self, obj): @property def parent_header(self): try: - return self._parent_header.get() + thread_id, parent_header = self._parent_header.get() except LookupError: return self._parent_header_global + if thread_id != threading.get_ident(): + return self._parent_header_global + return parent_header @parent_header.setter def parent_header(self, value): - self._parent_header.set(value) + self._parent_header.set((threading.get_ident(), 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)) + return self._parent_header.set((threading.get_ident(), extract_header(parent))) def reset_thread_parent(self, token): """Reset the parent header to undo the set_thread_parent call that returned the token.""" @@ -91,14 +94,14 @@ class ZMQShellDisplayHook(DisplayHook): session = Instance(Session, allow_none=True) pub_socket = Any(allow_none=True) - _parent_header: ContextVar[dict[str, Any]] + _parent_header: ContextVar[tuple[int, dict[str, Any]]] _thread_local = Any() msg: dict[str, t.Any] | None def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._parent_header = ContextVar("parent_header") - self._parent_header.set({}) + self._parent_header.set((threading.get_ident(), {})) self._parent_header_global = {} @default("_thread_local") @@ -131,18 +134,21 @@ def unregister_hook(self, hook): @property def parent_header(self): try: - return self._parent_header.get() + thread_id, parent_header = self._parent_header.get() except LookupError: return self._parent_header_global + if thread_id != threading.get_ident(): + return self._parent_header_global + return parent_header @parent_header.setter def parent_header(self, value): - self._parent_header.set(value) + self._parent_header.set((threading.get_ident(), 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)) + return self._parent_header.set((threading.get_ident(), extract_header(parent))) def reset_thread_parent(self, token): """Reset the parent header to undo the set_thread_parent call that returned the token.""" diff --git a/ipykernel/iostream.py b/ipykernel/iostream.py index ef6aa1994..b7379c416 100644 --- a/ipykernel/iostream.py +++ b/ipykernel/iostream.py @@ -539,10 +539,10 @@ def __init__( self.pub_thread = pub_thread self.name = name self.topic = b"stream." + name.encode() - self._parent_header: contextvars.ContextVar[dict[str, Any]] = contextvars.ContextVar( - "parent_header" + self._parent_header: contextvars.ContextVar[tuple[int, dict[str, Any]]] = ( + contextvars.ContextVar("parent_header") ) - self._parent_header.set({}) + self._parent_header.set((threading.get_ident(), {})) self._parent_header_global = {} self._master_pid = os.getpid() self._flush_pending = False @@ -598,14 +598,19 @@ def __init__( def parent_header(self): try: # asyncio or thread-specific - return self._parent_header.get() + thread_id, parent_header = self._parent_header.get() except LookupError: # global (fallback) return self._parent_header_global + if thread_id != threading.get_ident(): + # Contexts can be inherited by a new thread, but the parent is + # thread-specific unless explicitly set in that thread. + return self._parent_header_global + return parent_header @parent_header.setter def parent_header(self, value): - self._parent_header.set(value) + self._parent_header.set((threading.get_ident(), value)) self._parent_header_global = value def isatty(self): @@ -634,7 +639,7 @@ def _is_master_process(self): 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)) + return self._parent_header.set((threading.get_ident(), extract_header(parent))) def reset_thread_parent(self, token): """Reset the parent header to undo the set_thread_parent call that returned the token.""" diff --git a/ipykernel/zmqshell.py b/ipykernel/zmqshell.py index dd5216246..460d301d0 100644 --- a/ipykernel/zmqshell.py +++ b/ipykernel/zmqshell.py @@ -57,7 +57,7 @@ class ZMQDisplayPublisher(DisplayPublisher): session = Instance(Session, allow_none=True) pub_socket = Any(allow_none=True) - _parent_header: contextvars.ContextVar[dict[str, Any]] + _parent_header: contextvars.ContextVar[tuple[int, dict[str, Any]]] topic = CBytes(b"display_data") store_display_history = Bool( @@ -73,24 +73,27 @@ class ZMQDisplayPublisher(DisplayPublisher): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._parent_header = contextvars.ContextVar("parent_header") - self._parent_header.set({}) + self._parent_header.set((threading.get_ident(), {})) self._parent_header_global = {} @property def parent_header(self): try: - return self._parent_header.get() + thread_id, parent_header = self._parent_header.get() except LookupError: return self._parent_header_global + if thread_id != threading.get_ident(): + return self._parent_header_global + return parent_header @parent_header.setter def parent_header(self, value): - self._parent_header.set(value) + self._parent_header.set((threading.get_ident(), 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)) + return self._parent_header.set((threading.get_ident(), extract_header(parent))) def reset_thread_parent(self, token): """Reset the parent header to undo the set_thread_parent call that returned the token.""" @@ -546,10 +549,10 @@ 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[dict[str, typing.Any]] = contextvars.ContextVar( - "parent_header" + self._parent_header: contextvars.ContextVar[tuple[int, dict[str, typing.Any]]] = ( + contextvars.ContextVar("parent_header") ) - self._parent_header.set({}) + self._parent_header.set((threading.get_ident(), {})) self._parent_header_global = {} displayhook_class = Type(ZMQShellDisplayHook) @@ -732,13 +735,16 @@ def set_next_input(self, text, replace=False): @property def parent_header(self): try: - return self._parent_header.get() + thread_id, parent_header = self._parent_header.get() except LookupError: return self._parent_header_global + if thread_id != threading.get_ident(): + return self._parent_header_global + return parent_header @parent_header.setter def parent_header(self, value): - self._parent_header.set(value) + self._parent_header.set((threading.get_ident(), value)) self._parent_header_global = value def set_parent(self, parent): @@ -764,7 +770,12 @@ 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))] + tokens = [ + ( + self._parent_header.reset, + self._parent_header.set((threading.get_ident(), parent)), + ) + ] objs = [self.displayhook, self.display_pub, sys.stdout, sys.stderr] if hasattr(self, "_data_pub"): objs.append(self.data_pub) diff --git a/tests/test_kernel.py b/tests/test_kernel.py index fb9a7ad1f..13c5f68fd 100644 --- a/tests/test_kernel.py +++ b/tests/test_kernel.py @@ -131,6 +131,7 @@ def test_print_to_correct_cell_from_thread(explicit_parent: str): get_ipython().set_thread_parent sets the thread-local parent for only the thread. """ code = f"""\ + from contextvars import copy_context from threading import Event, Thread from time import sleep from IPython.display import display @@ -162,7 +163,8 @@ def thread_target(): print("after", flush=True) display(3) - thread = Thread(target=thread_target) + # Free-threaded Python inherits the caller's context by default. + thread = Thread(target=copy_context().run, args=(thread_target,)) thread.start() """ outputs = {}