feat(jupyter): support confirm and prompt in notebooks (#23592)

Closes: #22633 This commit adds support for `confirm` and `prompt` APIs, that instead of reading from stdin are using notebook frontend to show modal boxes and wait for answers. --------- Co-authored-by: Bartek Iwańczuk <[email protected]>
denoland · Jul 4, 2024 · f00f0f9 · f00f0f9
1 parent 96b527b
commit f00f0f9
Show file tree

Hide file tree

Showing 5 changed files with 165 additions and 13 deletions.
diff --git a/cli/js/40_jupyter.js b/cli/js/40_jupyter.js
@@ -337,7 +337,14 @@ async function formatInner(obj, raw) {
 internals.jupyter = { formatInner };
 
 function enableJupyter() {
-  const { op_jupyter_broadcast } = core.ops;
+  const { op_jupyter_broadcast, op_jupyter_input } = core.ops;
+
+  function input(
+    prompt,
+    password,
+  ) {
+    return op_jupyter_input(prompt, password);
+  }
 
   async function broadcast(
     msgType,
@@ -412,6 +419,45 @@ function enableJupyter() {
     return;
   }
 
+  /**
+   * Prompt for user confirmation (in Jupyter Notebook context)
+   * Override confirm and prompt because they depend on a tty
+   * and in the Deno.jupyter environment that doesn't exist.
+   * @param {string} message - The message to display.
+   * @returns {Promise<boolean>} User confirmation.
+   */
+  function confirm(message = "Confirm") {
+    const answer = input(`${message} [y/N] `, false);
+    return answer === "Y" || answer === "y";
+  }
+
+  /**
+   * Prompt for user input (in Jupyter Notebook context)
+   * @param {string} message - The message to display.
+   * @param {string} defaultValue - The value used if none is provided.
+   * @param {object} options Options
+   * @param {boolean} options.password Hide the output characters
+   * @returns {Promise<string>} The user input.
+   */
+  function prompt(
+    message = "Prompt",
+    defaultValue = "",
+    { password = false } = {},
+  ) {
+    if (defaultValue != "") {
+      message += ` [${defaultValue}]`;
+    }
+    const answer = input(`${message}`, password);
+
+    if (answer === "") {
+      return defaultValue;
+    }
+
+    return answer;
+  }
+
+  globalThis.confirm = confirm;
+  globalThis.prompt = prompt;
   globalThis.Deno.jupyter = {
     broadcast,
     display,

diff --git a/cli/ops/jupyter.rs b/cli/ops/jupyter.rs
@@ -1,24 +1,32 @@
 // Copyright 2018-2024 the Deno authors. All rights reserved. MIT license.
 
+// NOTE(bartlomieju): unfortunately it appears that clippy is broken
+// and can't allow a single line ignore for `await_holding_lock`.
+#![allow(clippy::await_holding_lock)]
+
 use std::cell::RefCell;
 use std::rc::Rc;
 use std::sync::Arc;
 
+use jupyter_runtime::InputRequest;
 use jupyter_runtime::JupyterMessage;
 use jupyter_runtime::JupyterMessageContent;
 use jupyter_runtime::KernelIoPubConnection;
 use jupyter_runtime::StreamContent;
 
 use deno_core::error::AnyError;
 use deno_core::op2;
+use deno_core::parking_lot::Mutex;
 use deno_core::serde_json;
 use deno_core::OpState;
 use tokio::sync::mpsc;
-use tokio::sync::Mutex;
+
+use crate::tools::jupyter::server::StdinConnectionProxy;
 
 deno_core::extension!(deno_jupyter,
   ops = [
     op_jupyter_broadcast,
+    op_jupyter_input,
   ],
   options = {
     sender: mpsc::UnboundedSender<StreamContent>,
@@ -32,6 +40,63 @@ deno_core::extension!(deno_jupyter,
   },
 );
 
+#[op2]
+#[string]
+pub fn op_jupyter_input(
+  state: &mut OpState,
+  #[string] prompt: String,
+  is_password: bool,
+) -> Result<Option<String>, AnyError> {
+  let (last_execution_request, stdin_connection_proxy) = {
+    (
+      state.borrow::<Arc<Mutex<Option<JupyterMessage>>>>().clone(),
+      state.borrow::<Arc<Mutex<StdinConnectionProxy>>>().clone(),
+    )
+  };
+
+  let maybe_last_request = last_execution_request.lock().clone();
+  if let Some(last_request) = maybe_last_request {
+    let JupyterMessageContent::ExecuteRequest(msg) = &last_request.content
+    else {
+      return Ok(None);
+    };
+
+    if !msg.allow_stdin {
+      return Ok(None);
+    }
+
+    let msg = JupyterMessage::new(
+      InputRequest {
+        prompt,
+        password: is_password,
+      }
+      .into(),
+      Some(&last_request),
+    );
+
+    let Ok(()) = stdin_connection_proxy.lock().tx.send(msg) else {
+      return Ok(None);
+    };
+
+    // Need to spawn a separate thread here, because `blocking_recv()` can't
+    // be used from the Tokio runtime context.
+    let join_handle = std::thread::spawn(move || {
+      stdin_connection_proxy.lock().rx.blocking_recv()
+    });
+    let Ok(Some(response)) = join_handle.join() else {
+      return Ok(None);
+    };
+
+    let JupyterMessageContent::InputReply(msg) = response.content else {
+      return Ok(None);
+    };
+
+    return Ok(Some(msg.value));
+  }
+
+  Ok(None)
+}
+
 #[op2(async)]
 pub async fn op_jupyter_broadcast(
   state: Rc<RefCell<OpState>>,
@@ -49,7 +114,7 @@ pub async fn op_jupyter_broadcast(
     )
   };
 
-  let maybe_last_request = last_execution_request.lock().await.clone();
+  let maybe_last_request = last_execution_request.lock().clone();
   if let Some(last_request) = maybe_last_request {
     let content = JupyterMessageContent::from_type_and_content(
       &message_type,
@@ -69,9 +134,7 @@ pub async fn op_jupyter_broadcast(
       .with_metadata(metadata)
       .with_buffers(buffers.into_iter().map(|b| b.to_vec().into()).collect());
 
-    (iopub_connection.lock().await)
-      .send(jupyter_message)
-      .await?;
+    iopub_connection.lock().send(jupyter_message).await?;
   }
 
   Ok(())

diff --git a/cli/tools/jupyter/mod.rs b/cli/tools/jupyter/mod.rs
@@ -179,6 +179,7 @@ pub async fn kernel(
     let mut op_state = op_state_rc.borrow_mut();
     op_state.put(startup_data.iopub_connection.clone());
     op_state.put(startup_data.last_execution_request.clone());
+    op_state.put(startup_data.stdin_connection_proxy.clone());
   }
 
   repl_session_proxy.start().await;

diff --git a/cli/tools/jupyter/server.rs b/cli/tools/jupyter/server.rs
@@ -3,6 +3,10 @@
 // This file is forked/ported from <https://github.com/evcxr/evcxr>
 // Copyright 2020 The Evcxr Authors. MIT license.
 
+// NOTE(bartlomieju): unfortunately it appears that clippy is broken
+// and can't allow a single line ignore for `await_holding_lock`.
+#![allow(clippy::await_holding_lock)]
+
 use std::collections::HashMap;
 use std::rc::Rc;
 use std::sync::Arc;
@@ -12,12 +16,12 @@ use crate::tools::repl;
 use deno_core::anyhow::bail;
 use deno_core::error::AnyError;
 use deno_core::futures;
+use deno_core::parking_lot::Mutex;
 use deno_core::serde_json;
 use deno_core::CancelFuture;
 use deno_core::CancelHandle;
 use tokio::sync::mpsc;
 use tokio::sync::oneshot;
-use tokio::sync::Mutex;
 
 use jupyter_runtime::messaging;
 use jupyter_runtime::AsChildOf;
@@ -40,8 +44,14 @@ pub struct JupyterServer {
   repl_session_proxy: JupyterReplProxy,
 }
 
+pub struct StdinConnectionProxy {
+  pub tx: mpsc::UnboundedSender<JupyterMessage>,
+  pub rx: mpsc::UnboundedReceiver<JupyterMessage>,
+}
+
 pub struct StartupData {
   pub iopub_connection: Arc<Mutex<KernelIoPubConnection>>,
+  pub stdin_connection_proxy: Arc<Mutex<StdinConnectionProxy>>,
   pub last_execution_request: Arc<Mutex<Option<JupyterMessage>>>,
 }
 
@@ -58,17 +68,27 @@ impl JupyterServer {
       connection_info.create_kernel_shell_connection().await?;
     let control_connection =
       connection_info.create_kernel_control_connection().await?;
-    let _stdin_connection =
+    let mut stdin_connection =
       connection_info.create_kernel_stdin_connection().await?;
     let iopub_connection =
       connection_info.create_kernel_iopub_connection().await?;
 
     let iopub_connection = Arc::new(Mutex::new(iopub_connection));
     let last_execution_request = Arc::new(Mutex::new(None));
 
+    let (stdin_tx1, mut stdin_rx1) =
+      mpsc::unbounded_channel::<JupyterMessage>();
+    let (stdin_tx2, stdin_rx2) = mpsc::unbounded_channel::<JupyterMessage>();
+
+    let stdin_connection_proxy = Arc::new(Mutex::new(StdinConnectionProxy {
+      tx: stdin_tx1,
+      rx: stdin_rx2,
+    }));
+
     let Ok(()) = setup_tx.send(StartupData {
       iopub_connection: iopub_connection.clone(),
       last_execution_request: last_execution_request.clone(),
+      stdin_connection_proxy,
     }) else {
       bail!("Failed to send startup data");
     };
@@ -82,6 +102,24 @@ impl JupyterServer {
       repl_session_proxy,
     };
 
+    let stdin_fut = deno_core::unsync::spawn(async move {
+      loop {
+        let Some(msg) = stdin_rx1.recv().await else {
+          return;
+        };
+        let Ok(()) = stdin_connection.send(msg).await else {
+          return;
+        };
+
+        let Ok(msg) = stdin_connection.read().await else {
+          return;
+        };
+        let Ok(()) = stdin_tx2.send(msg) else {
+          return;
+        };
+      }
+    });
+
     let hearbeat_fut = deno_core::unsync::spawn(async move {
       loop {
         if let Err(err) = heartbeat.single_heartbeat().await {
@@ -134,6 +172,7 @@ impl JupyterServer {
       shell_fut,
       stdio_fut,
       repl_session_fut,
+      stdin_fut,
     ]);
 
     if let Ok(result) = join_fut.or_cancel(cancel_handle).await {
@@ -148,13 +187,15 @@ impl JupyterServer {
     last_execution_request: Arc<Mutex<Option<JupyterMessage>>>,
     stdio_msg: StreamContent,
   ) {
-    let maybe_exec_result = last_execution_request.lock().await.clone();
+    let maybe_exec_result = last_execution_request.lock().clone();
     let Some(exec_request) = maybe_exec_result else {
       return;
     };
 
-    let mut iopub_conn = iopub_connection.lock().await;
-    let result = iopub_conn.send(stdio_msg.as_child_of(&exec_request)).await;
+    let result = iopub_connection
+      .lock()
+      .send(stdio_msg.as_child_of(&exec_request))
+      .await;
 
     if let Err(err) = result {
       log::error!("Output error: {}", err);
@@ -429,7 +470,7 @@ impl JupyterServer {
     if !execute_request.silent && execute_request.store_history {
       self.execution_count += 1;
     }
-    *self.last_execution_request.lock().await = Some(parent_message.clone());
+    *self.last_execution_request.lock() = Some(parent_message.clone());
 
     self
       .send_iopub(
@@ -613,7 +654,7 @@ impl JupyterServer {
     &mut self,
     message: JupyterMessage,
   ) -> Result<(), AnyError> {
-    self.iopub_connection.lock().await.send(message).await
+    self.iopub_connection.lock().send(message).await
   }
 }
 

diff --git a/runtime/js/99_main.js b/runtime/js/99_main.js
@@ -579,6 +579,7 @@ const NOT_IMPORTED_OPS = [
 
   // Related to `Deno.jupyter` API
   "op_jupyter_broadcast",
+  "op_jupyter_input",
 
   // Related to `Deno.test()` API
   "op_test_event_step_result_failed",