Replace _IOPubThread with BaseThread

ipykernel/inprocess/socket.py

@@ -65,4 +65,8 @@ async def poll(self, timeout=0):
         return statistics.current_buffer_used != 0
 
     def close(self):
-        pass
+        if self.is_shell:
+            self.in_send_stream.close()
+            self.in_receive_stream.close()
+            self.out_send_stream.close()
+            self.out_receive_stream.close()

ipykernel/iostream.py

@@ -24,6 +24,8 @@
 from anyio import create_task_group, run, sleep, to_thread
 from jupyter_client.session import extract_header
 
+from .thread import BaseThread
+
 # -----------------------------------------------------------------------------
 # Globals
 # -----------------------------------------------------------------------------
@@ -38,38 +40,6 @@
 # -----------------------------------------------------------------------------
 
 
-class _IOPubThread(Thread):
-    """A thread for a IOPub."""
-
-    def __init__(self, tasks, **kwargs):
-        """Initialize the thread."""
-        super().__init__(name="IOPub", **kwargs)
-        self._tasks = tasks
-        self.pydev_do_not_trace = True
-        self.is_pydev_daemon_thread = True
-        self.daemon = True
-        self.__stop = Event()
-
-    def run(self):
-        """Run the thread."""
-        self.name = "IOPub"
-        run(self._main)
-
-    async def _main(self):
-        async with create_task_group() as self._task_group:
-            for task in self._tasks:
-                self._task_group.start_soon(task)
-            await to_thread.run_sync(self.__stop.wait)
-            self._task_group.cancel_scope.cancel()
-
-    def stop(self):
-        """Stop the thread.
-
-        This method is threadsafe.
-        """
-        self.__stop.set()
-
-
 class IOPubThread:
     """An object for sending IOPub messages in a background thread
 
@@ -109,7 +79,9 @@ def __init__(self, socket: zmq_anyio.Socket, pipe=False):
         tasks = [self._handle_event, self._run_event_pipe_gc, self.socket.start]
         if pipe:
             tasks.append(self._handle_pipe_msgs)
-        self.thread = _IOPubThread(tasks)
+        self.thread = BaseThread(name="IOPub", daemon=True)
+        for task in tasks:
+            self.thread.start_soon(task)
 
     def _setup_event_pipe(self):
         """Create the PULL socket listening for events that should fire in this thread."""