From b276dfedd2d13de9918b9b87a3d543e7b5eca95e Mon Sep 17 00:00:00 2001 From: sigoden Date: Mon, 30 Oct 2023 10:07:01 +0800 Subject: feat: support textwrap (#171) * feat: support textwrap * improve cmd_render_stream * done cmd_render_stream * add `-w` alias to `--wrap` * done repl_render_stream * add `config.wrap_code` * remove cached config.wrap_width * fix unxpected duplicate lines on kitty * refactor markdown render * improve render * fix test --- src/render/cmd.rs | 56 +++++++--- src/render/markdown.rs | 272 ++++++++++++++++++++++++++++++++++++++----------- src/render/mod.rs | 25 ++--- src/render/repl.rs | 90 +++++++++------- 4 files changed, 319 insertions(+), 124 deletions(-) (limited to 'src/render') diff --git a/src/render/cmd.rs b/src/render/cmd.rs index 527c194..fb45b0a 100644 --- a/src/render/cmd.rs +++ b/src/render/cmd.rs @@ -5,15 +5,16 @@ use crate::repl::{ReplyStreamEvent, SharedAbortSignal}; use anyhow::Result; use crossbeam::channel::Receiver; +use textwrap::core::display_width; #[allow(clippy::unnecessary_wraps, clippy::module_name_repetitions)] pub fn cmd_render_stream( rx: &Receiver, - light_theme: bool, + render: &mut MarkdownRender, abort: &SharedAbortSignal, ) -> Result<()> { let mut buffer = String::new(); - let mut markdown_render = MarkdownRender::new(light_theme); + let mut col = 0; loop { if abort.aborted() { return Ok(()); @@ -23,28 +24,39 @@ pub fn cmd_render_stream( ReplyStreamEvent::Text(text) => { if text.contains('\n') { let text = format!("{buffer}{text}"); - let mut lines: Vec<&str> = text.split('\n').collect(); - buffer = lines.pop().unwrap_or_default().to_string(); - let output = lines.join("\n"); - print_now!("{}\n", markdown_render.render_block(&output)); + 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..]); + col = 0; } else { buffer = format!("{buffer}{text}"); - if !(markdown_render.is_code_block() - || buffer.len() < 60 + if !(render.is_code() + || buffer.len() < 40 || buffer.starts_with('#') || buffer.starts_with('>') || buffer.starts_with('|')) { - if let Some((output, remain)) = split_line(&buffer) { - print_now!("{}", markdown_render.render_line(&output)); + if let Some((head, remain)) = split_line_sematic(&buffer) { buffer = remain; + let input = format!("{}{head}", spaces(col)); + let output = render.render(&input); + let output = &output[col..]; + let (_, tail) = split_line_tail(output); + if output.contains('\n') { + col = display_width(tail); + } else { + col += display_width(output); + } + print_now!("{}", output); } } } } ReplyStreamEvent::Done => { - let output = markdown_render.render_block(&buffer); - print_now!("{}\n", output.trim_end()); + let input = format!("{}{buffer}", spaces(col)); + print_now!("{}\n", render.render(&input)); break; } } @@ -53,9 +65,9 @@ pub fn cmd_render_stream( Ok(()) } -fn split_line(line: &str) -> Option<(String, String)> { +fn split_line_sematic(text: &str) -> Option<(String, String)> { let mut balance: Vec = Vec::new(); - let chars: Vec = line.chars().collect(); + let chars: Vec = text.chars().collect(); let mut index = 0; let len = chars.len(); while index < len - 1 { @@ -82,6 +94,18 @@ fn split_line(line: &str) -> Option<(String, String)> { None } +pub(crate) fn split_line_tail(text: &str) -> (&str, &str) { + if let Some((head, tail)) = text.rsplit_once('\n') { + (head, tail) + } else { + ("", text) + } +} + +fn spaces(n: usize) -> String { + " ".repeat(n) +} + #[derive(Debug, Clone, Copy, Eq, PartialEq)] enum Kind { ParentheseStart, @@ -176,12 +200,12 @@ mod tests { macro_rules! assert_split_line { ($a:literal, $b:literal, true) => { assert_eq!( - split_line(&format!("{}{}", $a, $b)), + split_line_sematic(&format!("{}{}", $a, $b)), Some(($a.into(), $b.into())) ); }; ($a:literal, $b:literal, false) => { - assert_eq!(split_line(&format!("{}{}", $a, $b)), None); + assert_eq!(split_line_sematic(&format!("{}{}", $a, $b)), None); }; } diff --git a/src/render/markdown.rs b/src/render/markdown.rs index 66bb601..4866f3d 100644 --- a/src/render/markdown.rs +++ b/src/render/markdown.rs @@ -1,4 +1,6 @@ +use anyhow::{anyhow, Context, Result}; use crossterm::style::{Color, Stylize}; +use crossterm::terminal; use lazy_static::lazy_static; use std::collections::HashMap; use syntect::highlighting::{Color as SyntectColor, FontStyle, Style, Theme}; @@ -23,115 +25,175 @@ lazy_static! { #[allow(clippy::module_name_repetitions)] pub struct MarkdownRender { + options: RenderOptions, syntax_set: SyntaxSet, - md_theme: Theme, - code_color: Color, + md_theme: Option, + code_color: Option, md_syntax: SyntaxReference, code_syntax: Option, prev_line_type: LineType, + wrap_width: Option, } impl MarkdownRender { - pub fn new(light_theme: bool) -> Self { - let syntax_set: SyntaxSet = - bincode::deserialize_from(SYNTAXES).expect("invalid syntaxes binary"); - let md_theme: Theme = if light_theme { - bincode::deserialize_from(MD_THEME_LIGHT).expect("invalid theme binary") - } else { - bincode::deserialize_from(MD_THEME).expect("invalid theme binary") + pub fn init(options: RenderOptions) -> Result { + let syntax_set: SyntaxSet = bincode::deserialize_from(SYNTAXES) + .with_context(|| "MarkdownRender: invalid syntaxes binary")?; + + let md_theme: Option = match (options.highlight, options.light_theme) { + (false, _) => None, + (true, false) => Some( + bincode::deserialize_from(MD_THEME) + .with_context(|| "MarkdownRender: invalid theme binary")?, + ), + (true, true) => Some( + bincode::deserialize_from(MD_THEME_LIGHT) + .expect("MarkdownRender: invalid theme binary"), + ), }; - let code_color = get_code_color(&md_theme); + let code_color = md_theme.as_ref().map(get_code_color); let md_syntax = syntax_set.find_syntax_by_extension("md").unwrap().clone(); let line_type = LineType::Normal; - Self { + let wrap_width = match options.wrap.as_deref() { + None => None, + Some(value) => match terminal::size() { + Ok((columns, _)) => { + if value == "auto" { + Some(columns) + } else { + let value = value + .parse::() + .map_err(|_| anyhow!("Invalid wrap value"))?; + Some(columns.min(value)) + } + } + Err(_) => None, + }, + }; + Ok(Self { syntax_set, md_theme, code_color, md_syntax, code_syntax: None, prev_line_type: line_type, - } + wrap_width, + options, + }) + } + + pub(crate) const fn is_code(&self) -> bool { + matches!( + self.prev_line_type, + LineType::CodeBegin | LineType::CodeInner + ) } - pub fn render_block(&mut self, src: &str) -> String { - src.split('\n') - .map(|line| { - self.render_line_impl(line) - .unwrap_or_else(|| line.to_string()) - }) + pub fn render(&mut self, text: &str) -> String { + text.split('\n') + .map(|line| self.render_line_mut(line)) .collect::>() .join("\n") } pub fn render_line(&self, line: &str) -> String { - let output = if self.is_code_block() && detect_code_block(line).is_none() { - self.render_code_line(line) + let (_, code_syntax, is_code) = self.check_line(line); + if is_code { + self.highlint_code_line(line, &code_syntax) } else { - self.render_line_inner(line, &self.md_syntax) - }; - output.unwrap_or_else(|| line.to_string()) + self.highligh_line(line, &self.md_syntax, false) + } } - pub const fn is_code_block(&self) -> bool { - matches!( - self.prev_line_type, - LineType::CodeBegin | LineType::CodeInner - ) + fn render_line_mut(&mut self, line: &str) -> String { + let (line_type, code_syntax, is_code) = self.check_line(line); + let output = if is_code { + self.highlint_code_line(line, &code_syntax) + } else { + self.highligh_line(line, &self.md_syntax, false) + }; + self.prev_line_type = line_type; + self.code_syntax = code_syntax; + output } - fn render_line_impl(&mut self, line: &str) -> Option { + fn check_line(&self, line: &str) -> (LineType, Option, bool) { + let mut line_type = self.prev_line_type; + let mut code_syntax = self.code_syntax.clone(); + let mut is_code = false; if let Some(lang) = detect_code_block(line) { - match self.prev_line_type { + match line_type { LineType::Normal | LineType::CodeEnd => { - self.prev_line_type = LineType::CodeBegin; - self.code_syntax = if lang.is_empty() { + line_type = LineType::CodeBegin; + code_syntax = if lang.is_empty() { None } else { self.find_syntax(&lang).cloned() }; } LineType::CodeBegin | LineType::CodeInner => { - self.prev_line_type = LineType::CodeEnd; - self.code_syntax = None; + line_type = LineType::CodeEnd; + code_syntax = None; } } - self.render_line_inner(line, &self.md_syntax) } else { - match self.prev_line_type { - LineType::Normal => self.render_line_inner(line, &self.md_syntax), + match line_type { + LineType::Normal => {} LineType::CodeEnd => { - self.prev_line_type = LineType::Normal; - self.render_line_inner(line, &self.md_syntax) + line_type = LineType::Normal; } LineType::CodeBegin => { - if self.code_syntax.is_none() { + if code_syntax.is_none() { if let Some(syntax) = self.syntax_set.find_syntax_by_first_line(line) { - self.code_syntax = Some(syntax.clone()); + code_syntax = Some(syntax.clone()); } } - self.prev_line_type = LineType::CodeInner; - self.render_code_line(line) + line_type = LineType::CodeInner; + is_code = true; + } + LineType::CodeInner => { + is_code = true; } - LineType::CodeInner => self.render_code_line(line), } } + (line_type, code_syntax, is_code) } - fn render_line_inner(&self, line: &str, syntax: &SyntaxReference) -> Option { + fn highligh_line(&self, line: &str, syntax: &SyntaxReference, is_code: bool) -> String { let ws: String = line.chars().take_while(|c| c.is_whitespace()).collect(); - let trimed_line = &line[ws.len()..]; - let mut highlighter = HighlightLines::new(syntax, &self.md_theme); - let ranges = highlighter - .highlight_line(trimed_line, &self.syntax_set) - .ok()?; - Some(format!("{ws}{}", as_terminal_escaped(&ranges))) - } - - fn render_code_line(&self, line: &str) -> Option { - self.code_syntax.as_ref().map_or_else( - || Some(format!("{}", line.with(self.code_color))), - |syntax| self.render_line_inner(line, syntax), - ) + let trimed_line: &str = &line[ws.len()..]; + let mut line_highlighted = None; + if let Some(theme) = &self.md_theme { + let mut highlighter = HighlightLines::new(syntax, theme); + if let Ok(ranges) = highlighter.highlight_line(trimed_line, &self.syntax_set) { + line_highlighted = Some(format!("{ws}{}", as_terminal_escaped(&ranges))) + } + } + let line = line_highlighted.unwrap_or_else(|| line.into()); + self.wrap_line(line, is_code) + } + + fn highlint_code_line(&self, line: &str, code_syntax: &Option) -> String { + if let Some(syntax) = code_syntax { + self.highligh_line(line, syntax, true) + } else { + let line = match self.code_color { + Some(color) => line.with(color).to_string(), + None => line.to_string(), + }; + self.wrap_line(line, true) + } + } + + fn wrap_line(&self, line: String, is_code: bool) -> String { + if let Some(width) = self.wrap_width { + if is_code && !self.options.wrap_code { + return line; + } + textwrap::wrap(&line, width as usize).join("\n") + } else { + line + } } fn find_syntax(&self, lang: &str) -> Option<&SyntaxReference> { @@ -146,6 +208,30 @@ impl MarkdownRender { } } +#[derive(Debug, Clone, Default)] +pub struct RenderOptions { + pub highlight: bool, + pub light_theme: bool, + pub wrap: Option, + pub wrap_code: bool, +} + +impl RenderOptions { + pub(crate) fn new( + highlight: bool, + light_theme: bool, + wrap: Option, + wrap_code: bool, + ) -> Self { + Self { + highlight, + light_theme, + wrap, + wrap_code, + } + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum LineType { Normal, @@ -222,6 +308,44 @@ fn get_code_color(theme: &Theme) -> Color { mod tests { use super::*; + const TEXT: &str = r#" +To unzip a file in Rust, you can use the `zip` crate. Here's an example code that shows how to unzip a file: + +```rust +use std::fs::File; + +fn unzip_file(path: &str, output_dir: &str) -> Result<(), Box> { + todo!() +} +``` +"#; + const TEXT_NO_WRAP_CODE: &str = r#" +To unzip a file in Rust, you can use the `zip` crate. Here's an example code +that shows how to unzip a file: + +```rust +use std::fs::File; + +fn unzip_file(path: &str, output_dir: &str) -> Result<(), Box> { + todo!() +} +``` +"#; + + const TEXT_WRAP_ALL: &str = r#" +To unzip a file in Rust, you can use the `zip` crate. Here's an example code +that shows how to unzip a file: + +```rust +use std::fs::File; + +fn unzip_file(path: &str, output_dir: &str) -> Result<(), Box> { + todo!() +} +``` +"#; + #[test] fn test_assets() { let syntax_set: SyntaxSet = @@ -233,7 +357,37 @@ mod tests { #[test] fn test_render() { - let render = MarkdownRender::new(true); + let options = RenderOptions::default(); + let render = MarkdownRender::init(options).unwrap(); assert!(render.find_syntax("csharp").is_some()); } + + #[test] + fn no_theme() { + let options = RenderOptions::default(); + let mut render = MarkdownRender::init(options).unwrap(); + let output = render.render(TEXT); + assert_eq!(TEXT, output); + } + + #[test] + fn no_wrap_code() { + let options = RenderOptions::default(); + let mut render = MarkdownRender::init(options).unwrap(); + render.wrap_width = Some(80); + let output = render.render(TEXT); + assert_eq!(TEXT_NO_WRAP_CODE, output); + } + + #[test] + fn wrap_all() { + let options = RenderOptions { + wrap_code: true, + ..Default::default() + }; + let mut render = MarkdownRender::init(options).unwrap(); + render.wrap_width = Some(80); + let output = render.render(TEXT); + assert_eq!(TEXT_WRAP_ALL, output); + } } diff --git a/src/render/mod.rs b/src/render/mod.rs index e7c6228..9d7b06a 100644 --- a/src/render/mod.rs +++ b/src/render/mod.rs @@ -4,7 +4,7 @@ mod repl; use self::cmd::cmd_render_stream; #[allow(clippy::module_name_repetitions)] -pub use self::markdown::MarkdownRender; +pub use self::markdown::{MarkdownRender, RenderOptions}; use self::repl::repl_render_stream; use crate::client::Client; @@ -26,26 +26,27 @@ pub fn render_stream( abort: SharedAbortSignal, wg: WaitGroup, ) -> Result { - let (highlight, light_theme) = config.read().get_render_options(); - let mut stream_handler = if highlight { + let render_options = config.read().get_render_options(); + let mut stream_handler = { let (tx, rx) = unbounded(); let abort_clone = abort.clone(); spawn(move || { - let err = if repl { - repl_render_stream(&rx, light_theme, &abort) - } else { - cmd_render_stream(&rx, light_theme, &abort) + let run = move || { + if repl { + let mut render = MarkdownRender::init(render_options)?; + repl_render_stream(&rx, &mut render, &abort) + } else { + let mut render = MarkdownRender::init(render_options)?; + cmd_render_stream(&rx, &mut render, &abort) + } }; - if let Err(err) = err { + if let Err(err) = run() { let err = format!("{err:?}"); print_now!("{}\n\n", err.trim()); } drop(wg); }); - ReplyStreamHandler::new(Some(tx), repl, abort_clone) - } else { - drop(wg); - ReplyStreamHandler::new(None, repl, abort) + ReplyStreamHandler::new(tx, abort_clone) }; client.send_message_streaming(input, &mut stream_handler)?; let buffer = stream_handler.get_buffer(); diff --git a/src/render/repl.rs b/src/render/repl.rs index eb50750..c411670 100644 --- a/src/render/repl.rs +++ b/src/render/repl.rs @@ -1,4 +1,4 @@ -use super::MarkdownRender; +use super::{cmd::split_line_tail, MarkdownRender}; use crate::repl::{ReplyStreamEvent, SharedAbortSignal}; @@ -14,18 +14,18 @@ use std::{ io::{self, Stdout, Write}, time::{Duration, Instant}, }; -use unicode_width::UnicodeWidthStr; +use textwrap::core::display_width; #[allow(clippy::module_name_repetitions)] pub fn repl_render_stream( rx: &Receiver, - light_theme: bool, + render: &mut MarkdownRender, abort: &SharedAbortSignal, ) -> Result<()> { enable_raw_mode()?; let mut stdout = io::stdout(); - let ret = repl_render_stream_inner(rx, light_theme, abort, &mut stdout); + let ret = repl_render_stream_inner(rx, render, abort, &mut stdout); disable_raw_mode()?; @@ -34,15 +34,16 @@ pub fn repl_render_stream( fn repl_render_stream_inner( rx: &Receiver, - light_theme: bool, + render: &mut MarkdownRender, abort: &SharedAbortSignal, writer: &mut Stdout, ) -> Result<()> { let mut last_tick = Instant::now(); let tick_rate = Duration::from_millis(50); let mut buffer = String::new(); - let mut markdown_render = MarkdownRender::new(light_theme); let columns = terminal::size()?.0; + + let mut clear_rows = 0; loop { if abort.aborted() { return Ok(()); @@ -51,44 +52,45 @@ fn repl_render_stream_inner( if let Ok(evt) = rx.try_recv() { match evt { ReplyStreamEvent::Text(text) => { - if !buffer.is_empty() { - let buffer_width = buffer.width() as u16; - let need_rows = (buffer_width + columns - 1) / columns; - let (col, row) = cursor::position()?; - - if row + 1 >= need_rows { - if col == 0 { - queue!(writer, cursor::MoveTo(0, row - need_rows))?; - } else { - queue!(writer, cursor::MoveTo(0, row + 1 - need_rows))?; - } - } else { - queue!( - writer, - terminal::ScrollUp(need_rows - 1 - row), - cursor::MoveTo(0, 0) - )?; - } + let (col, mut row) = cursor::position()?; + + // fix unxpected 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 >= clear_rows { + queue!(writer, cursor::MoveTo(0, row - clear_rows))?; + } else { + let scroll_rows = clear_rows - row - 1; + queue!( + writer, + terminal::ScrollUp(scroll_rows), + cursor::MoveTo(0, 0), + )?; } if text.contains('\n') { let text = format!("{buffer}{text}"); - let mut lines: Vec<&str> = text.split('\n').collect(); - buffer = lines.pop().unwrap_or_default().to_string(); - let output = markdown_render.render_block(&lines.join("\n")); - for line in output.split('\n') { - queue!( - writer, - style::Print(line), - style::Print("\n"), - cursor::MoveLeft(columns), - )?; - } + 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),)?; + clear_rows = 0; } else { buffer = format!("{buffer}{text}"); - let output = markdown_render.render_line(&buffer); - queue!(writer, style::Print(&output))?; + let output = render.render_line(&buffer); + if output.contains('\n') { + let (head, tail) = split_line_tail(&output); + clear_rows = print_block(writer, head, columns)?; + queue!(writer, style::Print(&tail),)?; + } else { + queue!(writer, style::Print(&output))?; + let buffer_width = display_width(&output) as u16; + let need_rows = (buffer_width + columns - 1) / columns; + clear_rows = need_rows.saturating_sub(1); + } } writer.flush()?; @@ -128,3 +130,17 @@ fn repl_render_stream_inner( } 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) +} -- cgit v1.2.3