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

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<Option<bool>> {
     Some(parse_bool(&value))
 }
 
-fn parse_bool(value: &str) -> Option<bool> {
-    match value {
-        "1" | "true" => Some(true),
-        "0" | "false" => Some(false),
-        _ => None,
-    }
-}
-
 fn complete_bool(value: bool) -> Vec<String> {
     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<bool> {
+    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()
 }