From b3162a72b50f88d78ca3d57c1f8b3421cecfe2b4 Mon Sep 17 00:00:00 2001 From: sigoden Date: Sun, 21 Apr 2024 06:52:22 +0800 Subject: refactor: shell execution with spinner and select prompt (#425) --- src/main.rs | 53 +++++++++++++++++++++++++++++++++------------------- src/render/stream.rs | 51 +------------------------------------------------- src/utils/mod.rs | 2 ++ src/utils/spinner.rs | 52 +++++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 89 insertions(+), 69 deletions(-) create mode 100644 src/utils/spinner.rs 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) -> Vec { 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(()) + } +} -- cgit v1.2.3