From 4bc74f7e9006fc7f182c9dc6ea309ea83c5bb4f9 Mon Sep 17 00:00:00 2001 From: sigoden Date: Tue, 8 Oct 2024 12:38:31 +0800 Subject: [PATCH] refactor: NO_COLOR, highlight and render_error (#915) - If NO_COLOR is enabled, AIChat will not output any ANSI color codes. - `config.highlight` will only affect rendering markdown - improve rendering error --- src/config/mod.rs | 14 ++------------ src/main.rs | 6 ++---- src/render/mod.rs | 11 +++-------- src/repl/highlighter.rs | 34 ++++++++++++++-------------------- src/repl/mod.rs | 2 +- src/utils/mod.rs | 21 +++++++++++++++++++-- 6 files changed, 41 insertions(+), 47 deletions(-) diff --git a/src/config/mod.rs b/src/config/mod.rs index 1c669a7..8e3e2c3 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -1969,10 +1969,8 @@ impl Config { if let Some(Some(v)) = read_env_bool("highlight") { self.highlight = v; } - if let Ok(value) = env::var("NO_COLOR") { - if let Some(false) = parse_bool(&value) { - self.highlight = false; - } + if *NO_COLOR { + self.highlight = false; } if let Some(Some(v)) = read_env_bool("light_theme") { self.light_theme = v; @@ -2166,14 +2164,6 @@ fn read_env_bool(key: &str) -> Option> { Some(parse_bool(&value)) } -fn parse_bool(value: &str) -> Option { - match value { - "1" | "true" => Some(true), - "0" | "false" => Some(false), - _ => None, - } -} - fn complete_bool(value: bool) -> Vec { vec![(!value).to_string()] } diff --git a/src/main.rs b/src/main.rs index 1476408..3c52f0d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -34,7 +34,7 @@ use parking_lot::RwLock; use simplelog::{format_description, ConfigBuilder, LevelFilter, SimpleLogger, WriteLogger}; use std::{ env, - io::{stderr, stdin, Read}, + io::{stdin, Read}, process, sync::Arc, }; @@ -54,10 +54,8 @@ async fn main() -> Result<()> { }; setup_logger(working_mode.is_serve())?; let config = Arc::new(RwLock::new(Config::init(working_mode)?)); - let highlight = config.read().highlight; if let Err(err) = run(config, cli, text).await { - let highlight = stderr().is_terminal() && highlight; - render_error(err, highlight); + render_error(err); std::process::exit(1); } Ok(()) diff --git a/src/render/mod.rs b/src/render/mod.rs index 4e55e9f..97161f8 100644 --- a/src/render/mod.rs +++ b/src/render/mod.rs @@ -4,7 +4,7 @@ mod stream; pub use self::markdown::{MarkdownRender, RenderOptions}; use self::stream::{markdown_stream, raw_stream}; -use crate::utils::{error_text, AbortSignal, IS_STDOUT_TERMINAL}; +use crate::utils::{error_text, pretty_error, AbortSignal, IS_STDOUT_TERMINAL}; use crate::{client::SseEvent, config::GlobalConfig}; use anyhow::Result; @@ -25,11 +25,6 @@ pub async fn render_stream( ret.map_err(|err| err.context("Failed to reader stream")) } -pub fn render_error(err: anyhow::Error, highlight: bool) { - let err = format!("Error: {err:?}"); - if highlight { - eprintln!("{}", error_text(&err)); - } else { - eprintln!("{err}"); - } +pub fn render_error(err: anyhow::Error) { + eprintln!("{}", error_text(&pretty_error(&err))); } diff --git a/src/repl/highlighter.rs b/src/repl/highlighter.rs index 65cb1cc..ab6174c 100644 --- a/src/repl/highlighter.rs +++ b/src/repl/highlighter.rs @@ -1,34 +1,28 @@ use super::REPL_COMMANDS; -use crate::config::GlobalConfig; +use crate::{config::GlobalConfig, utils::NO_COLOR}; use nu_ansi_term::{Color, Style}; use reedline::{Highlighter, StyledText}; -pub struct ReplHighlighter { - config: GlobalConfig, -} +const DEFAULT_COLOR: Color = Color::Default; +const MATCH_COLOR: Color = Color::Green; + +pub struct ReplHighlighter; impl ReplHighlighter { - pub fn new(config: &GlobalConfig) -> Self { - Self { - config: config.clone(), - } + pub fn new(_config: &GlobalConfig) -> Self { + Self } } impl Highlighter for ReplHighlighter { fn highlight(&self, line: &str, _cursor: usize) -> StyledText { - let color = Color::Default; - let match_color = if self.config.read().highlight { - Color::Green - } else { - color - }; - let mut styled_text = StyledText::new(); - if REPL_COMMANDS.iter().any(|cmd| line.contains(cmd.name)) { + if *NO_COLOR { + styled_text.push((Style::default(), line.to_string())); + } else if REPL_COMMANDS.iter().any(|cmd| line.contains(cmd.name)) { let matches: Vec<&str> = REPL_COMMANDS .iter() .filter(|cmd| line.contains(cmd.name)) @@ -43,11 +37,11 @@ impl Highlighter for ReplHighlighter { }); let buffer_split: Vec<&str> = line.splitn(2, &longest_match).collect(); - styled_text.push((Style::new().fg(color), buffer_split[0].to_string())); - styled_text.push((Style::new().fg(match_color), longest_match)); - styled_text.push((Style::new().fg(color), buffer_split[1].to_string())); + styled_text.push((Style::new().fg(DEFAULT_COLOR), buffer_split[0].to_string())); + styled_text.push((Style::new().fg(MATCH_COLOR), longest_match)); + styled_text.push((Style::new().fg(DEFAULT_COLOR), buffer_split[1].to_string())); } else { - styled_text.push((Style::new().fg(color), line.to_string())); + styled_text.push((Style::new().fg(DEFAULT_COLOR), line.to_string())); } styled_text diff --git a/src/repl/mod.rs b/src/repl/mod.rs index 5a99ae9..a7b2573 100644 --- a/src/repl/mod.rs +++ b/src/repl/mod.rs @@ -211,7 +211,7 @@ impl Repl { } } Err(err) => { - render_error(err, self.config.read().highlight); + render_error(err); println!() } } diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 81ae0ff..0bd1d41 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -31,6 +31,7 @@ use unicode_segmentation::UnicodeSegmentation; lazy_static::lazy_static! { pub static ref CODE_BLOCK_RE: Regex = Regex::new(r"(?ms)```\w*(.*)```").unwrap(); pub static ref IS_STDOUT_TERMINAL: bool = std::io::stdout().is_terminal(); + pub static ref NO_COLOR: bool = env::var("NO_COLOR").ok().and_then(|v| parse_bool(&v)).unwrap_or_default() || !*IS_STDOUT_TERMINAL; } pub fn now() -> String { @@ -46,6 +47,14 @@ pub fn normalize_env_name(value: &str) -> String { value.replace('-', "_").to_ascii_uppercase() } +pub fn parse_bool(value: &str) -> Option { + match value { + "1" | "true" => Some(true), + "0" | "false" => Some(false), + _ => None, + } +} + pub fn estimate_token_length(text: &str) -> usize { let words: Vec<&str> = text.unicode_words().collect(); let mut output: f32 = 0.0; @@ -127,9 +136,11 @@ pub fn fuzzy_match(text: &str, pattern: &str) -> bool { pub fn pretty_error(err: &anyhow::Error) -> String { let mut output = vec![]; - output.push(err.to_string()); - output.push("Caused by:".to_string()); + output.push(format!("Error: {err}")); for (i, cause) in err.chain().skip(1).enumerate() { + if i == 0 { + output.push("Caused by:".to_string()); + } output.push(format!(" {i}: {cause}")); } output.push(String::new()); @@ -145,6 +156,9 @@ pub fn warning_text(input: &str) -> String { } pub fn color_text(input: &str, color: nu_ansi_term::Color) -> String { + if *NO_COLOR { + return input.to_string(); + } nu_ansi_term::Style::new() .fg(color) .paint(input) @@ -152,6 +166,9 @@ pub fn color_text(input: &str, color: nu_ansi_term::Color) -> String { } pub fn dimmed_text(input: &str) -> String { + if *NO_COLOR { + return input.to_string(); + } nu_ansi_term::Style::new().dimmed().paint(input).to_string() }