diff options
| author | sigoden <sigoden@gmail.com> | 2023-11-02 09:53:54 +0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2023-11-02 09:53:54 +0800 |
| commit | 444f4ebe9de8aa68afdf8d39b734a8543b9b5c4b (patch) | |
| tree | d8b0e16faa0270fda2a6c7afa5fe18c1970b9a93 /src/render/mod.rs | |
| parent | 5c7bfd92ff3e557477969be9db0638ef0d3d3659 (diff) | |
| download | aichat-444f4ebe9de8aa68afdf8d39b734a8543b9b5c4b.tar.gz | |
refactor: improve code quanity (#196)
- rewrite Repl, remove ReplHandler
- move ReplyStreamHandler to repl/ and rename it to ReplyHandler
- deprecate utils::print_now
- refactor session info
Diffstat (limited to 'src/render/mod.rs')
| -rw-r--r-- | src/render/mod.rs | 84 |
1 files changed, 76 insertions, 8 deletions
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, +} |
