Let get_parent decide the channel to get parent header (#1128)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: dby-tmwctw

ipykernel/comm/comm.py

@@ -3,7 +3,6 @@
 # Copyright (c) IPython Development Team.
 # Distributed under the terms of the Modified BSD License.
 
-import threading
 import uuid
 from typing import Optional
 from warnings import warn
@@ -12,7 +11,6 @@
 import traitlets.config
 from traitlets import Bool, Bytes, Instance, Unicode, default
 
-from ipykernel.control import CONTROL_THREAD_NAME
 from ipykernel.jsonutil import json_clean
 from ipykernel.kernelbase import Kernel
 
@@ -32,11 +30,6 @@ def publish_msg(self, msg_type, data=None, metadata=None, buffers=None, **keys):
         metadata = {} if metadata is None else metadata
         content = json_clean(dict(data=data, comm_id=self.comm_id, **keys))
 
-        if threading.current_thread().name == CONTROL_THREAD_NAME:
-            channel_from_which_to_get_parent_header = "control"
-        else:
-            channel_from_which_to_get_parent_header = "shell"
-
         if self.kernel is None:
             self.kernel = Kernel.instance()
 
@@ -45,7 +38,7 @@ def publish_msg(self, msg_type, data=None, metadata=None, buffers=None, **keys):
             msg_type,
             content,
             metadata=json_clean(metadata),
-            parent=self.kernel.get_parent(channel_from_which_to_get_parent_header),
+            parent=self.kernel.get_parent(),
             ident=self.topic,
             buffers=buffers,
         )

ipykernel/ipkernel.py

@@ -675,7 +675,6 @@ def do_apply(self, content, bufs, msg_id, reply_metadata):
                 "error",
                 reply_content,
                 ident=self._topic("error"),
-                channel="shell",
             )
             self.log.info("Exception in apply request:\n%s", "\n".join(reply_content["traceback"]))
             result_buf = []

ipykernel/kernelbase.py

@@ -11,6 +11,7 @@
 import os
 import socket
 import sys
+import threading
 import time
 import typing as t
 import uuid
@@ -19,6 +20,8 @@
 from functools import partial
 from signal import SIGINT, SIGTERM, Signals, default_int_handler, signal
 
+from .control import CONTROL_THREAD_NAME
+
 if sys.platform != "win32":
     from signal import SIGKILL
 else:
@@ -187,7 +190,7 @@ def _parent_header(self):
             DeprecationWarning,
             stacklevel=2,
         )
-        return self.get_parent(channel="shell")
+        return self.get_parent()
 
     # Time to sleep after flushing the stdout/err buffers in each execute
     # cycle.  While this introduces a hard limit on the minimal latency of the
@@ -598,7 +601,7 @@ def _publish_debug_event(self, event):
             self.iopub_socket,
             "debug_event",
             event,
-            parent=self.get_parent("control"),
+            parent=self.get_parent(),
             ident=self._topic("debug_event"),
         )
 
@@ -614,7 +617,7 @@ def set_parent(self, ident, parent, channel="shell"):
         self._parent_ident[channel] = ident
         self._parents[channel] = parent
 
-    def get_parent(self, channel="shell"):
+    def get_parent(self, channel=None):
         """Get the parent request associated with a channel.
 
         .. versionadded:: 6
@@ -629,6 +632,14 @@ def get_parent(self, channel="shell"):
         message : dict
             the parent message for the most recent request on the channel.
         """
+
+        if channel is None:
+            # If a channel is not specified, get information from current thread
+            if threading.current_thread().name == CONTROL_THREAD_NAME:
+                channel = "control"
+            else:
+                channel = "shell"
+
         return self._parents.get(channel, {})
 
     def send_response(
@@ -641,7 +652,7 @@ def send_response(
         track=False,
         header=None,
         metadata=None,
-        channel="shell",
+        channel=None,
     ):
         """Send a response to the message we're currently processing.