From eb30d90391ef73429589a98f34dfe4d3d88d0a38 Mon Sep 17 00:00:00 2001 From: sigoden Date: Wed, 8 Nov 2023 18:26:38 +0800 Subject: refactor: improve render (#235) * refactor: redesign render - if stdout is not terminal, just write reply text to stdout - rename repl_render_stream to markdown_stream - deprecate cmd_render_stream - use raw_stream to just print streaming reply text * optimize rendering error * optimize render_stream --- src/render/cmd.rs | 64 ----------------- src/render/markdown.rs | 40 ----------- src/render/mod.rs | 43 +++++++----- src/render/repl.rs | 163 ------------------------------------------- src/render/stream.rs | 183 +++++++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 210 insertions(+), 283 deletions(-) delete mode 100644 src/render/cmd.rs delete mode 100644 src/render/repl.rs create mode 100644 src/render/stream.rs (limited to 'src/render') diff --git a/src/render/cmd.rs b/src/render/cmd.rs deleted file mode 100644 index e942025..0000000 --- a/src/render/cmd.rs +++ /dev/null @@ -1,64 +0,0 @@ -use super::{MarkdownRender, ReplyEvent}; - -use crate::utils::{split_line_sematic, split_line_tail, AbortSignal}; - -use anyhow::Result; -use crossbeam::channel::Receiver; -use textwrap::core::display_width; - -pub fn cmd_render_stream( - rx: &Receiver, - render: &mut MarkdownRender, - abort: &AbortSignal, -) -> Result<()> { - let mut buffer = String::new(); - let mut indent = 0; - loop { - if abort.aborted() { - return Ok(()); - } - if let Ok(evt) = rx.try_recv() { - match evt { - ReplyEvent::Text(text) => { - if text.contains('\n') { - let text = format!("{buffer}{text}"); - let (head, tail) = split_line_tail(&text); - buffer = tail.to_string(); - let output = render.render_with_indent(head, indent); - println!("{}", output); - indent = 0; - } else { - buffer = format!("{buffer}{text}"); - if !(render.is_code() - || buffer.len() < 40 - || buffer.starts_with('#') - || buffer.starts_with('>') - || buffer.starts_with('|')) - { - if let Some((head, remain)) = split_line_sematic(&buffer) { - buffer = remain; - let output = render.render_with_indent(&head, indent); - let (_, tail) = split_line_tail(&output); - if let Some(width) = render.wrap_width() { - if output.contains('\n') { - indent = display_width(tail); - } else { - indent += display_width(&output); - } - indent %= width as usize; - } - print!("{}", output); - } - } - } - } - ReplyEvent::Done => { - let output = render.render_with_indent(&buffer, indent); - println!("{}", output); - break; - } - } - } - } - Ok(()) -} diff --git a/src/render/markdown.rs b/src/render/markdown.rs index f300854..44e24b3 100644 --- a/src/render/markdown.rs +++ b/src/render/markdown.rs @@ -64,17 +64,6 @@ impl MarkdownRender { }) } - pub(crate) const fn is_code(&self) -> bool { - matches!( - self.prev_line_type, - LineType::CodeBegin | LineType::CodeInner - ) - } - - pub(crate) const fn wrap_width(&self) -> Option { - self.wrap_width - } - pub fn render(&mut self, text: &str) -> String { text.split('\n') .map(|line| self.render_line_mut(line)) @@ -82,16 +71,6 @@ impl MarkdownRender { .join("\n") } - pub fn render_with_indent(&mut self, text: &str, indent: usize) -> String { - let text = format!("{}{}", " ".repeat(indent), text); - let output = self.render(&text); - if output.starts_with('\n') { - output - } else { - output.chars().skip(indent).collect() - } - } - pub fn render_line(&self, line: &str) -> String { let (_, code_syntax, is_code) = self.check_line(line); if is_code { @@ -377,23 +356,4 @@ std::error::Error>> { let output = render.render(TEXT); assert_eq!(TEXT_WRAP_ALL, output); } - - #[test] - fn wrap_with_indent() { - let options = RenderOptions::default(); - let mut render = MarkdownRender::init(options).unwrap(); - render.wrap_width = Some(80); - - let input = "To unzip a file in Rust, you can use the `zip` crate. Here's an example code"; - let output = render.render_with_indent(input, 40); - let expect = - "To unzip a file in Rust, you can use the\n`zip` crate. Here's an example code"; - assert_eq!(output, expect); - - let input = "Unzip a file"; - let output = render.render_with_indent(input, 76); - let expect = "\nUnzip a file"; - - assert_eq!(output, expect); - } } diff --git a/src/render/mod.rs b/src/render/mod.rs index 2b97557..1e20ccc 100644 --- a/src/render/mod.rs +++ b/src/render/mod.rs @@ -1,10 +1,8 @@ -mod cmd; mod markdown; -mod repl; +mod stream; -use self::cmd::cmd_render_stream; pub use self::markdown::{MarkdownRender, RenderOptions}; -use self::repl::repl_render_stream; +use self::stream::{markdown_stream, raw_stream}; use crate::client::Client; use crate::config::GlobalConfig; @@ -13,17 +11,19 @@ use crate::utils::AbortSignal; use anyhow::{Context, Result}; use crossbeam::channel::{unbounded, Sender}; use crossbeam::sync::WaitGroup; +use is_terminal::IsTerminal; use nu_ansi_term::{Color, Style}; +use std::io::stdout; use std::thread::spawn; pub fn render_stream( input: &str, client: &dyn Client, config: &GlobalConfig, - repl: bool, abort: AbortSignal, - wg: WaitGroup, ) -> Result { + let wg = WaitGroup::new(); + let wg_cloned = wg.clone(); let render_options = config.read().get_render_options()?; let mut stream_handler = { let (tx, rx) = unbounded(); @@ -31,33 +31,44 @@ pub fn render_stream( let highlight = config.read().highlight; spawn(move || { let run = move || { - if repl { + if stdout().is_terminal() { let mut render = MarkdownRender::init(render_options)?; - repl_render_stream(&rx, &mut render, &abort) + markdown_stream(&rx, &mut render, &abort) } else { - let mut render = MarkdownRender::init(render_options)?; - cmd_render_stream(&rx, &mut render, &abort) + raw_stream(&rx, &abort) } }; if let Err(err) = run() { render_error(err, highlight); } - drop(wg); + drop(wg_cloned); }); ReplyHandler::new(tx, abort_clone) }; - client.send_message_streaming(input, &mut stream_handler)?; - let buffer = stream_handler.get_buffer(); - Ok(buffer.to_string()) + let ret = client.send_message_streaming(input, &mut stream_handler); + wg.wait(); + let output = stream_handler.get_buffer().to_string(); + match ret { + Ok(_) => { + println!(); + Ok(output) + } + Err(err) => { + if !output.is_empty() { + println!(); + } + Err(err) + } + } } pub fn render_error(err: anyhow::Error, highlight: bool) { let err = format!("{err:?}"); if highlight { let style = Style::new().fg(Color::Red); - println!("{}", style.paint(err.trim())); + eprintln!("{}", style.paint(err)); } else { - println!("{}", err.trim()); + eprintln!("{err}"); } } diff --git a/src/render/repl.rs b/src/render/repl.rs deleted file mode 100644 index eeb5420..0000000 --- a/src/render/repl.rs +++ /dev/null @@ -1,163 +0,0 @@ -use super::{MarkdownRender, ReplyEvent}; - -use crate::utils::{split_line_tail, AbortSignal}; - -use anyhow::Result; -use crossbeam::channel::Receiver; -use crossterm::{ - cursor, - event::{self, Event, KeyCode, KeyModifiers}, - queue, style, - terminal::{self, disable_raw_mode, enable_raw_mode}, -}; -use std::{ - io::{self, Stdout, Write}, - time::{Duration, Instant}, -}; -use textwrap::core::display_width; - -pub fn repl_render_stream( - rx: &Receiver, - render: &mut MarkdownRender, - abort: &AbortSignal, -) -> Result<()> { - enable_raw_mode()?; - let mut stdout = io::stdout(); - - let ret = repl_render_stream_inner(rx, render, abort, &mut stdout); - - disable_raw_mode()?; - - ret -} - -fn repl_render_stream_inner( - rx: &Receiver, - render: &mut MarkdownRender, - abort: &AbortSignal, - writer: &mut Stdout, -) -> Result<()> { - let mut last_tick = Instant::now(); - let tick_rate = Duration::from_millis(50); - - let mut buffer = String::new(); - let mut buffer_rows = 1; - - let columns = terminal::size()?.0; - - loop { - if abort.aborted() { - return Ok(()); - } - - if let Ok(evt) = rx.try_recv() { - match evt { - ReplyEvent::Text(text) => { - let (col, mut row) = cursor::position()?; - - // Fix unexpected duplicate lines on kitty, see https://github.com/sigoden/aichat/issues/105 - if col == 0 && row > 0 && display_width(&buffer) == columns as usize { - row -= 1; - } - - if row + 1 >= buffer_rows { - queue!(writer, cursor::MoveTo(0, row + 1 - buffer_rows),)?; - } else { - let scroll_rows = buffer_rows - row - 1; - queue!( - writer, - terminal::ScrollUp(scroll_rows), - cursor::MoveTo(0, 0), - )?; - } - - // No guarantee that text returned by render will not be re-layouted, so it is better to clear it. - queue!(writer, terminal::Clear(terminal::ClearType::FromCursorDown))?; - - if text.contains('\n') { - let text = format!("{buffer}{text}"); - let (head, tail) = split_line_tail(&text); - buffer = tail.to_string(); - let output = render.render(head); - print_block(writer, &output, columns)?; - queue!(writer, style::Print(&buffer),)?; - - // No guarantee the buffer width of the buffer will not exceed the number of columns. - // So we calculate the number of rows needed, rather than setting it directly to 1. - buffer_rows = need_rows(&buffer, columns); - } else { - buffer = format!("{buffer}{text}"); - let output = render.render_line(&buffer); - if output.contains('\n') { - let (head, tail) = split_line_tail(&output); - buffer_rows = print_block(writer, head, columns)?; - queue!(writer, style::Print(&tail),)?; - - // Same as above. - buffer_rows += need_rows(tail, columns); - } else { - queue!(writer, style::Print(&output))?; - buffer_rows = need_rows(&output, columns); - } - } - - writer.flush()?; - } - ReplyEvent::Done => { - #[cfg(target_os = "windows")] - let eol = "\n\n"; - #[cfg(not(target_os = "windows"))] - let eol = "\n"; - queue!(writer, style::Print(eol))?; - writer.flush()?; - - break; - } - } - continue; - } - - let timeout = tick_rate - .checked_sub(last_tick.elapsed()) - .unwrap_or_else(|| Duration::from_secs(0)); - if crossterm::event::poll(timeout)? { - if let Event::Key(key) = event::read()? { - match key.code { - KeyCode::Char('c') if key.modifiers == KeyModifiers::CONTROL => { - abort.set_ctrlc(); - return Ok(()); - } - KeyCode::Char('d') if key.modifiers == KeyModifiers::CONTROL => { - abort.set_ctrld(); - return Ok(()); - } - _ => {} - } - } - } - - if last_tick.elapsed() >= tick_rate { - last_tick = Instant::now(); - } - } - Ok(()) -} - -fn print_block(writer: &mut Stdout, text: &str, columns: u16) -> Result { - let mut num = 0; - for line in text.split('\n') { - queue!( - writer, - style::Print(line), - style::Print("\n"), - cursor::MoveLeft(columns), - )?; - num += 1; - } - Ok(num) -} - -fn need_rows(text: &str, columns: u16) -> u16 { - let buffer_width = display_width(text).max(1) as u16; - (buffer_width + columns - 1) / columns -} diff --git a/src/render/stream.rs b/src/render/stream.rs new file mode 100644 index 0000000..189862d --- /dev/null +++ b/src/render/stream.rs @@ -0,0 +1,183 @@ +use super::{MarkdownRender, ReplyEvent}; + +use crate::utils::AbortSignal; + +use anyhow::Result; +use crossbeam::channel::Receiver; +use crossterm::{ + cursor, + event::{self, Event, KeyCode, KeyModifiers}, + queue, style, + terminal::{self, disable_raw_mode, enable_raw_mode}, +}; +use std::{ + io::{self, Stdout, Write}, + time::{Duration, Instant}, +}; +use textwrap::core::display_width; + +pub fn markdown_stream( + rx: &Receiver, + render: &mut MarkdownRender, + abort: &AbortSignal, +) -> Result<()> { + enable_raw_mode()?; + let mut stdout = io::stdout(); + + let ret = markdown_stream_inner(rx, render, abort, &mut stdout); + + disable_raw_mode()?; + + ret +} + +pub fn raw_stream(rx: &Receiver, abort: &AbortSignal) -> Result<()> { + loop { + if abort.aborted() { + return Ok(()); + } + if let Ok(evt) = rx.try_recv() { + match evt { + ReplyEvent::Text(text) => { + print!("{}", text); + } + ReplyEvent::Done => { + break; + } + } + } + } + Ok(()) +} + +fn markdown_stream_inner( + rx: &Receiver, + render: &mut MarkdownRender, + abort: &AbortSignal, + writer: &mut Stdout, +) -> Result<()> { + let mut last_tick = Instant::now(); + let tick_rate = Duration::from_millis(50); + + let mut buffer = String::new(); + let mut buffer_rows = 1; + + let columns = terminal::size()?.0; + + loop { + if abort.aborted() { + return Ok(()); + } + + if let Ok(evt) = rx.try_recv() { + match evt { + ReplyEvent::Text(text) => { + let (col, mut row) = cursor::position()?; + + // Fix unexpected duplicate lines on kitty, see https://github.com/sigoden/aichat/issues/105 + if col == 0 && row > 0 && display_width(&buffer) == columns as usize { + row -= 1; + } + + if row + 1 >= buffer_rows { + queue!(writer, cursor::MoveTo(0, row + 1 - buffer_rows),)?; + } else { + let scroll_rows = buffer_rows - row - 1; + queue!( + writer, + terminal::ScrollUp(scroll_rows), + cursor::MoveTo(0, 0), + )?; + } + + // No guarantee that text returned by render will not be re-layouted, so it is better to clear it. + queue!(writer, terminal::Clear(terminal::ClearType::FromCursorDown))?; + + if text.contains('\n') { + let text = format!("{buffer}{text}"); + let (head, tail) = split_line_tail(&text); + buffer = tail.to_string(); + let output = render.render(head); + print_block(writer, &output, columns)?; + queue!(writer, style::Print(&buffer),)?; + + // No guarantee the buffer width of the buffer will not exceed the number of columns. + // So we calculate the number of rows needed, rather than setting it directly to 1. + buffer_rows = need_rows(&buffer, columns); + } else { + buffer = format!("{buffer}{text}"); + let output = render.render_line(&buffer); + if output.contains('\n') { + let (head, tail) = split_line_tail(&output); + buffer_rows = print_block(writer, head, columns)?; + queue!(writer, style::Print(&tail),)?; + + // Same as above. + buffer_rows += need_rows(tail, columns); + } else { + queue!(writer, style::Print(&output))?; + buffer_rows = need_rows(&output, columns); + } + } + + writer.flush()?; + } + ReplyEvent::Done => { + break; + } + } + continue; + } + + let timeout = tick_rate + .checked_sub(last_tick.elapsed()) + .unwrap_or_else(|| Duration::from_secs(0)); + if crossterm::event::poll(timeout)? { + if let Event::Key(key) = event::read()? { + match key.code { + KeyCode::Char('c') if key.modifiers == KeyModifiers::CONTROL => { + abort.set_ctrlc(); + return Ok(()); + } + KeyCode::Char('d') if key.modifiers == KeyModifiers::CONTROL => { + abort.set_ctrld(); + return Ok(()); + } + _ => {} + } + } + } + + if last_tick.elapsed() >= tick_rate { + last_tick = Instant::now(); + } + } + Ok(()) +} + +fn print_block(writer: &mut Stdout, text: &str, columns: u16) -> Result { + let mut num = 0; + for line in text.split('\n') { + queue!( + writer, + style::Print(line), + style::Print("\n"), + cursor::MoveLeft(columns), + )?; + num += 1; + } + Ok(num) +} + +fn split_line_tail(text: &str) -> (&str, &str) { + if let Some((head, tail)) = text.rsplit_once('\n') { + (head, tail) + } else { + ("", text) + } +} + +fn need_rows(text: &str, columns: u16) -> u16 { + let buffer_width = display_width(text).max(1) as u16; + (buffer_width + columns - 1) / columns +} -- cgit v1.2.3