diff options
Diffstat (limited to 'src/render')
| -rw-r--r-- | src/render/cmd.rs | 19 | ||||
| -rw-r--r-- | src/render/mod.rs | 84 | ||||
| -rw-r--r-- | src/render/repl.rs | 16 |
3 files changed, 93 insertions, 26 deletions
diff --git a/src/render/cmd.rs b/src/render/cmd.rs index 115170b..e09bd71 100644 --- a/src/render/cmd.rs +++ b/src/render/cmd.rs @@ -1,7 +1,6 @@ -use super::MarkdownRender; +use super::{MarkdownRender, ReplyEvent}; -use crate::print_now; -use crate::repl::{ReplyStreamEvent, SharedAbortSignal}; +use crate::repl::AbortSignal; use crate::utils::{spaces, split_line_sematic, split_line_tail}; use anyhow::Result; @@ -9,9 +8,9 @@ use crossbeam::channel::Receiver; use textwrap::core::display_width; pub fn cmd_render_stream( - rx: &Receiver<ReplyStreamEvent>, + rx: &Receiver<ReplyEvent>, render: &mut MarkdownRender, - abort: &SharedAbortSignal, + abort: &AbortSignal, ) -> Result<()> { let mut buffer = String::new(); let mut col = 0; @@ -21,14 +20,14 @@ pub fn cmd_render_stream( } if let Ok(evt) = rx.try_recv() { match evt { - ReplyStreamEvent::Text(text) => { + ReplyEvent::Text(text) => { if text.contains('\n') { let text = format!("{buffer}{text}"); let (head, tail) = split_line_tail(&text); buffer = tail.to_string(); let input = format!("{}{head}", spaces(col)); let output = render.render(&input); - print_now!("{}\n", &output[col..]); + println!("{}", &output[col..]); col = 0; } else { buffer = format!("{buffer}{text}"); @@ -51,14 +50,14 @@ pub fn cmd_render_stream( col += display_width(output); } } - print_now!("{}", output); + print!("{}", output); } } } } - ReplyStreamEvent::Done => { + ReplyEvent::Done => { let input = format!("{}{buffer}", spaces(col)); - print_now!("{}\n", render.render(&input)); + println!("{}", render.render(&input)); break; } } diff --git a/src/render/mod.rs b/src/render/mod.rs index 03b67ea..27f9e81 100644 --- a/src/render/mod.rs +++ b/src/render/mod.rs @@ -8,12 +8,12 @@ use self::repl::repl_render_stream; use crate::client::Client; use crate::config::SharedConfig; -use crate::print_now; -use crate::repl::{ReplyStreamHandler, SharedAbortSignal}; +use crate::repl::AbortSignal; -use anyhow::Result; -use crossbeam::channel::unbounded; +use anyhow::{Context, Result}; +use crossbeam::channel::{unbounded, Sender}; use crossbeam::sync::WaitGroup; +use nu_ansi_term::{Color, Style}; use std::thread::spawn; pub fn render_stream( @@ -21,13 +21,14 @@ pub fn render_stream( client: &dyn Client, config: &SharedConfig, repl: bool, - abort: SharedAbortSignal, + abort: AbortSignal, wg: WaitGroup, ) -> Result<String> { let render_options = config.read().get_render_options()?; let mut stream_handler = { let (tx, rx) = unbounded(); let abort_clone = abort.clone(); + let highlight = config.read().highlight; spawn(move || { let run = move || { if repl { @@ -39,14 +40,81 @@ pub fn render_stream( } }; if let Err(err) = run() { - let err = format!("{err:?}"); - print_now!("\n{}\n\n", err.trim()); + render_error(err, highlight); } drop(wg); }); - ReplyStreamHandler::new(tx, abort_clone) + ReplyHandler::new(tx, abort_clone) }; client.send_message_streaming(input, &mut stream_handler)?; let buffer = stream_handler.get_buffer(); Ok(buffer.to_string()) } + +pub fn render_error(err: anyhow::Error, highlight: bool) { + let err = format!("{err:?}\n"); + if highlight { + let style = Style::new().fg(Color::Red); + println!("{}", style.paint(err.trim())); + } else { + println!("{}", err.trim()); + } +} + +pub struct ReplyHandler { + sender: Sender<ReplyEvent>, + buffer: String, + abort: AbortSignal, +} + +impl ReplyHandler { + pub fn new(sender: Sender<ReplyEvent>, abort: AbortSignal) -> Self { + Self { + sender, + abort, + buffer: String::new(), + } + } + + pub fn text(&mut self, text: &str) -> Result<()> { + if self.buffer.is_empty() && text == "\n\n" { + return Ok(()); + } + self.buffer.push_str(text); + let ret = self + .sender + .send(ReplyEvent::Text(text.to_string())) + .with_context(|| "Failed to send ReplyEvent:Text"); + self.safe_ret(ret)?; + Ok(()) + } + + pub fn done(&mut self) -> Result<()> { + let ret = self + .sender + .send(ReplyEvent::Done) + .with_context(|| "Failed to send ReplyEvent::Done"); + self.safe_ret(ret)?; + Ok(()) + } + + pub fn get_buffer(&self) -> &str { + &self.buffer + } + + pub fn get_abort(&self) -> AbortSignal { + self.abort.clone() + } + + fn safe_ret(&self, ret: Result<()>) -> Result<()> { + if ret.is_err() && self.abort.aborted() { + return Ok(()); + } + ret + } +} + +pub enum ReplyEvent { + Text(String), + Done, +} diff --git a/src/render/repl.rs b/src/render/repl.rs index 7363923..efad0e7 100644 --- a/src/render/repl.rs +++ b/src/render/repl.rs @@ -1,6 +1,6 @@ -use super::MarkdownRender; +use super::{MarkdownRender, ReplyEvent}; -use crate::repl::{ReplyStreamEvent, SharedAbortSignal}; +use crate::repl::AbortSignal; use crate::utils::split_line_tail; use anyhow::Result; @@ -18,9 +18,9 @@ use std::{ use textwrap::core::display_width; pub fn repl_render_stream( - rx: &Receiver<ReplyStreamEvent>, + rx: &Receiver<ReplyEvent>, render: &mut MarkdownRender, - abort: &SharedAbortSignal, + abort: &AbortSignal, ) -> Result<()> { enable_raw_mode()?; let mut stdout = io::stdout(); @@ -33,9 +33,9 @@ pub fn repl_render_stream( } fn repl_render_stream_inner( - rx: &Receiver<ReplyStreamEvent>, + rx: &Receiver<ReplyEvent>, render: &mut MarkdownRender, - abort: &SharedAbortSignal, + abort: &AbortSignal, writer: &mut Stdout, ) -> Result<()> { let mut last_tick = Instant::now(); @@ -51,7 +51,7 @@ fn repl_render_stream_inner( if let Ok(evt) = rx.try_recv() { match evt { - ReplyStreamEvent::Text(text) => { + ReplyEvent::Text(text) => { let (col, mut row) = cursor::position()?; // fix unexpected duplicate lines on kitty, see https://github.com/sigoden/aichat/issues/105 @@ -95,7 +95,7 @@ fn repl_render_stream_inner( writer.flush()?; } - ReplyStreamEvent::Done => { + ReplyEvent::Done => { #[cfg(target_os = "windows")] let eol = "\n\n"; #[cfg(not(target_os = "windows"))] |