From 3f8dd20374ecb21120b043624391cd4acda865f2 Mon Sep 17 00:00:00 2001
From: Min RK <benjaminrk@gmail.com>
Date: Mon, 23 Sep 2024 09:44:55 +0200
Subject: [PATCH] call wait_for_ready() in runapp

ensures kernel is ready before running

avoids lost output if iopub isn't connected yet

and shutdown kernel when finished, rather than relying on the kernel shutting itself down
---
 jupyter_client/runapp.py | 29 +++++++----------------------
 1 file changed, 7 insertions(+), 22 deletions(-)

diff --git a/jupyter_client/runapp.py b/jupyter_client/runapp.py
index 9ed4b154..bea5d5cf 100644
--- a/jupyter_client/runapp.py
+++ b/jupyter_client/runapp.py
@@ -3,10 +3,9 @@
 # Distributed under the terms of the Modified BSD License.
 from __future__ import annotations
 
-import queue
+import atexit
 import signal
 import sys
-import time
 import typing as t
 
 from jupyter_core.application import JupyterApp, base_aliases, base_flags
@@ -73,7 +72,8 @@ def initialize(self, argv: list[str] | None = None) -> None:  # type:ignore[over
         super().initialize(argv)
         JupyterConsoleApp.initialize(self)
         signal.signal(signal.SIGINT, self.handle_sigint)
-        self.init_kernel_info()
+        if self.kernel_manager:
+            atexit.register(self.kernel_manager.shutdown_kernel)
 
     def handle_sigint(self, *args: t.Any) -> None:
         """Handle SIGINT."""
@@ -82,28 +82,11 @@ def handle_sigint(self, *args: t.Any) -> None:
         else:
             self.log.error("Cannot interrupt kernels we didn't start.\n")
 
-    def init_kernel_info(self) -> None:
-        """Wait for a kernel to be ready, and store kernel info"""
-        timeout = self.kernel_timeout
-        tic = time.time()
-        self.kernel_client.hb_channel.unpause()
-        msg_id = self.kernel_client.kernel_info()
-        while True:
-            try:
-                reply = self.kernel_client.get_shell_msg(timeout=1)
-            except queue.Empty as e:
-                if (time.time() - tic) > timeout:
-                    msg = "Kernel didn't respond to kernel_info_request"
-                    raise RuntimeError(msg) from e
-            else:
-                if reply["parent_header"].get("msg_id") == msg_id:
-                    self.kernel_info = reply["content"]
-                    return
-
     def start(self) -> None:
         """Start the application."""
         self.log.debug("jupyter run: starting...")
         super().start()
+        self.kernel_client.wait_for_ready(timeout=self.kernel_timeout)
         if self.filenames_to_run:
             for filename in self.filenames_to_run:
                 self.log.debug("jupyter run: executing `%s`", filename)
@@ -112,8 +95,10 @@ def start(self) -> None:
                     reply = self.kernel_client.execute_interactive(code, timeout=OUTPUT_TIMEOUT)
                     return_code = 0 if reply["content"]["status"] == "ok" else 1
                     if return_code:
-                        raise Exception("jupyter-run error running '%s'" % filename)
+                        msg = f"jupyter-run error running '{filename}'"
+                        raise Exception(msg)
         else:
+            self.log.debug("jupyter run: executing from stdin")
             code = sys.stdin.read()
             reply = self.kernel_client.execute_interactive(code, timeout=OUTPUT_TIMEOUT)
             return_code = 0 if reply["content"]["status"] == "ok" else 1