Use UnknownSubshellError for mising subshell

Author: krassowski

ipykernel/kernelbase.py

@@ -59,6 +59,7 @@
 
 from ._version import kernel_protocol_version
 from .iostream import OutStream
+from .subshell_manager import UnknownSubshellError
 from .utils import LazyDict, _async_in_context
 
 psutil: t.Any | None = None
@@ -603,8 +604,8 @@ async def shell_channel_thread_main(self, msg):
                 subshell_manager = self.shell_channel_thread.manager
                 try:
                     socket = subshell_manager.get_shell_channel_to_subshell_socket(subshell_id)
-                except KeyError:
-                    self._send_unknown_subshell_reply(idents, msg3, subshell_id)
+                except UnknownSubshellError as err:
+                    self._send_unknown_subshell_reply(idents, msg3, err)
                     return
                 assert socket is not None
                 socket.send_multipart(msg, copy=False)
@@ -1382,7 +1383,7 @@ def _send_abort_reply(self, stream, msg, idents):
             ident=idents,
         )
 
-    def _send_unknown_subshell_reply(self, idents, msg, subshell_id) -> None:
+    def _send_unknown_subshell_reply(self, idents, msg, err: UnknownSubshellError) -> None:
         """Send an error reply to a request addressed to a subshell that is not there.
 
         Runs in the shell channel thread, so it writes to the shell socket
@@ -1396,17 +1397,12 @@ def _send_unknown_subshell_reply(self, idents, msg, subshell_id) -> None:
         if not self.session:
             return
         msg_type = msg["header"]["msg_type"]
-        self.log.warning(
-            "Cannot handle %s %s: unknown subshell_id %r",
-            msg_type,
-            msg["header"]["msg_id"],
-            subshell_id,
-        )
+        self.log.warning("Cannot handle %s %s: %s", msg_type, msg["header"]["msg_id"], err)
         self._publish_status("busy", "shell", parent=msg)
         content = {
             "status": "error",
-            "ename": "KeyError",
-            "evalue": f"Unknown subshell_id {subshell_id!r}",
+            "ename": type(err).__name__,
+            "evalue": str(err),
             "traceback": [],
         }
         md = self.init_metadata(msg)

ipykernel/subshell_manager.py

@@ -19,6 +19,20 @@
 from .utils import _async_in_context
 
 
+class UnknownSubshellError(KeyError):
+    """A subshell_id that does not name an existing subshell.
+
+    .. versionadded:: 7.4
+    """
+
+    def __init__(self, subshell_id: str) -> None:
+        super().__init__(subshell_id)
+        self.subshell_id = subshell_id
+
+    def __str__(self) -> str:
+        return f"Unknown subshell_id {self.subshell_id!r}"
+
+
 class SubshellManager:
     """A manager of subshells.
 
@@ -87,13 +101,24 @@ def close(self) -> None:
         self._main_to_shell_channel.close()
         self._shell_channel_to_main.close()
 
+    def _get_subshell(self, subshell_id: str) -> SubshellThread:
+        """Return the thread of the specified subshell.
+
+        The caller must hold ``_lock_cache``.  Raises ``UnknownSubshellError`` if
+        there is no such subshell.
+        """
+        try:
+            return self._cache[subshell_id]
+        except KeyError:
+            raise UnknownSubshellError(subshell_id) from None
+
     def get_shell_channel_to_subshell_pair(self, subshell_id: str | None) -> SocketPair:
         """Return the inproc socket pair used to send messages from the shell channel
         to a particular subshell or main shell."""
         if subshell_id is None:
             return self._shell_channel_to_main
         with self._lock_cache:
-            return self._cache[subshell_id].shell_channel_to_subshell
+            return self._get_subshell(subshell_id).shell_channel_to_subshell
 
     def get_subshell_to_shell_channel_socket(self, subshell_id: str | None) -> zmq.Socket[t.Any]:
         """Return the socket used by a particular subshell or main shell to send
@@ -102,7 +127,7 @@ def get_subshell_to_shell_channel_socket(self, subshell_id: str | None) -> zmq.S
         if subshell_id is None:
             return self._main_to_shell_channel.from_socket
         with self._lock_cache:
-            return self._cache[subshell_id].subshell_to_shell_channel.from_socket
+            return self._get_subshell(subshell_id).subshell_to_shell_channel.from_socket
 
     def get_shell_channel_to_subshell_socket(self, subshell_id: str | None) -> zmq.Socket[t.Any]:
         """Return the socket used by the shell channel to send messages to a particular
@@ -113,12 +138,12 @@ def get_shell_channel_to_subshell_socket(self, subshell_id: str | None) -> zmq.S
     def get_subshell_aborting(self, subshell_id: str) -> bool:
         """Get the boolean aborting flag of the specified subshell."""
         with self._lock_cache:
-            return self._cache[subshell_id].aborting
+            return self._get_subshell(subshell_id).aborting
 
     def get_subshell_asyncio_lock(self, subshell_id: str) -> asyncio.Lock:
         """Return the asyncio lock belonging to the specified subshell."""
         with self._lock_cache:
-            return self._cache[subshell_id].asyncio_lock
+            return self._get_subshell(subshell_id).asyncio_lock
 
     def list_subshell(self) -> list[str]:
         """Return list of current subshell ids.
@@ -141,7 +166,7 @@ def set_on_recv_callback(self, on_recv_callback):
     def set_subshell_aborting(self, subshell_id: str, aborting: bool) -> None:
         """Set the aborting flag of the specified subshell."""
         with self._lock_cache:
-            self._cache[subshell_id].aborting = aborting
+            self._get_subshell(subshell_id).aborting = aborting
 
     def subshell_id_from_thread_id(self, thread_id: int) -> str | None:
         """Return subshell_id of the specified thread_id.
@@ -185,14 +210,15 @@ def _create_subshell(self) -> str:
     def _delete_subshell(self, subshell_id: str) -> None:
         """Delete subshell identified by subshell_id.
 
-        Raises key error if subshell_id not in cache.
+        Raises ``UnknownSubshellError`` if subshell_id not in cache.
         """
         assert current_thread().name == SHELL_CHANNEL_THREAD_NAME
 
         with self._lock_cache:
-            subshell_threwad = self._cache.pop(subshell_id)
+            subshell_thread = self._get_subshell(subshell_id)
+            del self._cache[subshell_id]
 
-        self._stop_subshell(subshell_threwad)
+        self._stop_subshell(subshell_thread)
 
     def _process_control_request(
         self,

tests/test_subshells.py

@@ -399,14 +399,20 @@ def test_unknown_subshell_id():
     with new_kernel() as kc:
         subshell_id = create_subshell_helper(kc)["subshell_id"]
         delete_subshell_helper(kc, subshell_id)
+
+        # Deleting it again names the missing subshell in the same way.
+        content = delete_subshell_helper(kc, subshell_id)
+        assert content["status"] == "error"
+        assert content["evalue"] == f"Unknown subshell_id {subshell_id!r}"
         flush_channels(kc)
 
         msg = execute_request(kc, "a = 1", subshell_id)
         msg_id = msg["header"]["msg_id"]
 
         reply = get_reply(kc, msg_id, TIMEOUT)
         assert reply["content"]["status"] == "error"
-        assert subshell_id in reply["content"]["evalue"]
+        assert reply["content"]["ename"] == "UnknownSubshellError"
+        assert reply["content"]["evalue"] == f"Unknown subshell_id {subshell_id!r}"
 
         states = []
         while True: