refactor: shell execution with spinner and select prompt (#425)

4 files changed, 89 insertions, 69 deletions

diff --git a/src/main.rs b/src/main.rs
index 66492da..600381b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -20,13 +20,14 @@ use crate::utils::{
 
 use anyhow::{bail, Result};
 use clap::Parser;
-use inquire::validator::Validation;
-use inquire::Text;
+use inquire::{Select, Text};
 use is_terminal::IsTerminal;
 use parking_lot::RwLock;
 use std::io::{stderr, stdin, stdout, Read};
-use std::process;
-use std::sync::Arc;
+use std::sync::{mpsc, Arc};
+use std::time::Duration;
+use std::{process, thread};
+use utils::Spinner;
 
 fn main() -> Result<()> {
     let cli = Cli::parse();
@@ -153,7 +154,11 @@ fn start_interactive(config: &GlobalConfig) -> Result<()> {
 fn execute(config: &GlobalConfig, mut input: Input) -> Result<()> {
     let client = init_client(config)?;
     config.read().maybe_print_send_tokens(&input);
-    let mut eval_str = client.send_message(input.clone())?;
+    let (tx, rx) = mpsc::sync_channel::<()>(0);
+    thread::spawn(move || run_spinner(rx));
+    let ret = client.send_message(input.clone());
+    tx.send(())?;
+    let mut eval_str = ret?;
     if let Ok(true) = CODE_BLOCK_RE.is_match(&eval_str) {
         eval_str = extract_block(&eval_str);
     }
@@ -166,27 +171,22 @@ fn execute(config: &GlobalConfig, mut input: Input) -> Result<()> {
         return Ok(());
     }
     if stdout().is_terminal() {
-        println!("{}", markdown_render.render(&eval_str).trim());
         let mut explain = false;
         loop {
-            let answer = Text::new("[1]:execute [2]:explain [3]:revise [4]:cancel")
-                .with_default("1")
-                .with_validator(|input: &str| match matches!(input, "1" | "2" | "3" | "4") {
-                    true => Ok(Validation::Valid),
-                    false => Ok(Validation::Invalid(
-                        "Select a number between 1 and 4.".into(),
-                    )),
-                })
-                .prompt()?;
+            let answer = Select::new(
+                markdown_render.render(&eval_str).trim(),
+                vec!["✅ Execute", "📙 Explain", "🤔 Revise", "❌ Cancel"],
+            )
+            .prompt()?;
 
-            match answer.as_str() {
-                "1" => {
+            match answer {
+                "✅ Execute" => {
                     let code = run_command(&eval_str)?;
                     if code != 0 {
                         process::exit(code);
                     }
                 }
-                "2" => {
+                "📙 Explain" => {
                     if !explain {
                         config.write().set_role(EXPLAIN_ROLE)?;
                     }
@@ -196,7 +196,7 @@ fn execute(config: &GlobalConfig, mut input: Input) -> Result<()> {
                     explain = true;
                     continue;
                 }
-                "3" => {
+                "🤔 Revise" => {
                     let revision = Text::new("Enter your revision:").prompt()?;
                     let text = format!(
                         "[INST] {} [/INST]\n{eval_str}\n[INST] {revision} [/INST]\n",
@@ -246,3 +246,18 @@ fn create_input(
     };
     Ok(Some(input))
 }
+
+fn run_spinner(rx: mpsc::Receiver<()>) -> Result<()> {
+    let mut writer = stdout();
+    let mut spinner = Spinner::new(" Generating");
+    loop {
+        spinner.step(&mut writer)?;
+        if let Ok(()) = rx.try_recv() {
+            spinner.stop(&mut writer)?;
+            break;
+        }
+        thread::sleep(Duration::from_millis(50))
+    }
+    spinner.stop(&mut writer)?;
+    Ok(())
+}
diff --git a/src/render/stream.rs b/src/render/stream.rs
index 8bb4a0e..9fdad5a 100644
--- a/src/render/stream.rs
+++ b/src/render/stream.rs
@@ -1,6 +1,6 @@
 use super::{MarkdownRender, ReplyEvent};
 
-use crate::utils::AbortSignal;
+use crate::utils::{AbortSignal, Spinner};
 
 use anyhow::Result;
 use crossbeam::channel::Receiver;
@@ -164,55 +164,6 @@ fn markdown_stream_inner(
     Ok(())
 }
 
-struct Spinner {
-    index: usize,
-    message: String,
-    stopped: bool,
-}
-
-impl Spinner {
-    const DATA: [&'static str; 10] = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
-
-    fn new(message: &str) -> Self {
-        Spinner {
-            index: 0,
-            message: message.to_string(),
-            stopped: false,
-        }
-    }
-
-    fn step(&mut self, writer: &mut Stdout) -> Result<()> {
-        if self.stopped {
-            return Ok(());
-        }
-        let frame = Self::DATA[self.index % Self::DATA.len()];
-        let dots = ".".repeat((self.index / 5) % 4);
-        let line = format!("{frame}{}{:<3}", self.message, dots);
-        queue!(writer, cursor::MoveToColumn(0), style::Print(line),)?;
-        if self.index == 0 {
-            queue!(writer, cursor::Hide)?;
-        }
-        writer.flush()?;
-        self.index += 1;
-        Ok(())
-    }
-
-    fn stop(&mut self, writer: &mut Stdout) -> Result<()> {
-        if self.stopped {
-            return Ok(());
-        }
-        self.stopped = true;
-        queue!(
-            writer,
-            cursor::MoveToColumn(0),
-            terminal::Clear(terminal::ClearType::FromCursorDown),
-            cursor::Show
-        )?;
-        writer.flush()?;
-        Ok(())
-    }
-}
-
 fn gather_events(rx: &Receiver<ReplyEvent>) -> Vec<ReplyEvent> {
     let mut texts = vec![];
     let mut done = false;
diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index 15af4ca..3df1a09 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -2,12 +2,14 @@ mod abort_signal;
 mod clipboard;
 mod prompt_input;
 mod render_prompt;
+mod spinner;
 mod tiktoken;
 
 pub use self::abort_signal::{create_abort_signal, AbortSignal};
 pub use self::clipboard::set_text;
 pub use self::prompt_input::*;
 pub use self::render_prompt::render_prompt;
+pub use self::spinner::Spinner;
 pub use self::tiktoken::cl100k_base_singleton;
 
 use fancy_regex::Regex;
diff --git a/src/utils/spinner.rs b/src/utils/spinner.rs
new file mode 100644
index 0000000..c684901
--- /dev/null
+++ b/src/utils/spinner.rs
@@ -0,0 +1,52 @@
+use anyhow::Result;
+use crossterm::{cursor, queue, style, terminal};
+use std::io::{Stdout, Write};
+
+pub struct Spinner {
+    index: usize,
+    message: String,
+    stopped: bool,
+}
+
+impl Spinner {
+    const DATA: [&'static str; 10] = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
+
+    pub fn new(message: &str) -> Self {
+        Spinner {
+            index: 0,
+            message: message.to_string(),
+            stopped: false,
+        }
+    }
+
+    pub fn step(&mut self, writer: &mut Stdout) -> Result<()> {
+        if self.stopped {
+            return Ok(());
+        }
+        let frame = Self::DATA[self.index % Self::DATA.len()];
+        let dots = ".".repeat((self.index / 5) % 4);
+        let line = format!("{frame}{}{:<3}", self.message, dots);
+        queue!(writer, cursor::MoveToColumn(0), style::Print(line),)?;
+        if self.index == 0 {
+            queue!(writer, cursor::Hide)?;
+        }
+        writer.flush()?;
+        self.index += 1;
+        Ok(())
+    }
+
+    pub fn stop(&mut self, writer: &mut Stdout) -> Result<()> {
+        if self.stopped {
+            return Ok(());
+        }
+        self.stopped = true;
+        queue!(
+            writer,
+            cursor::MoveToColumn(0),
+            terminal::Clear(terminal::ClearType::FromCursorDown),
+            cursor::Show
+        )?;
+        writer.flush()?;
+        Ok(())
+    }
+}