vtcode_core/utils/
ansi.rs

1use crate::ui::theme;
2use anstream::{AutoStream, ColorChoice};
3use anstyle::{Reset, Style};
4use anstyle_query::{clicolor, clicolor_force, no_color, term_supports_color};
5use anyhow::Result;
6use std::io::{self, Write};
7
8/// Styles available for rendering messages
9#[derive(Clone, Copy)]
10pub enum MessageStyle {
11    Info,
12    Error,
13    Output,
14    Response,
15    Tool,
16    User,
17}
18
19impl MessageStyle {
20    fn style(self) -> Style {
21        let styles = theme::active_styles();
22        match self {
23            Self::Info => styles.info,
24            Self::Error => styles.error,
25            Self::Output => styles.output,
26            Self::Response => styles.response,
27            Self::Tool => styles.tool,
28            Self::User => styles.user,
29        }
30    }
31
32    fn indent(self) -> &'static str {
33        match self {
34            Self::Response | Self::Tool => "  ",
35            _ => "",
36        }
37    }
38}
39
40/// Renderer with deferred output buffering
41pub struct AnsiRenderer {
42    writer: AutoStream<io::Stdout>,
43    buffer: String,
44    color: bool,
45}
46
47impl AnsiRenderer {
48    /// Create a new renderer for stdout
49    pub fn stdout() -> Self {
50        let color =
51            clicolor_force() || (!no_color() && clicolor().unwrap_or_else(term_supports_color));
52        let choice = if color {
53            ColorChoice::Auto
54        } else {
55            ColorChoice::Never
56        };
57        Self {
58            writer: AutoStream::new(std::io::stdout(), choice),
59            buffer: String::new(),
60            color,
61        }
62    }
63
64    /// Push text into the buffer
65    pub fn push(&mut self, text: &str) {
66        self.buffer.push_str(text);
67    }
68
69    /// Flush the buffer with the given style
70    pub fn flush(&mut self, style: MessageStyle) -> Result<()> {
71        let style = style.style();
72        if self.color {
73            writeln!(self.writer, "{style}{}{Reset}", self.buffer)?;
74        } else {
75            writeln!(self.writer, "{}", self.buffer)?;
76        }
77        self.writer.flush()?;
78        self.buffer.clear();
79        Ok(())
80    }
81
82    /// Convenience for writing a single line
83    pub fn line(&mut self, style: MessageStyle, text: &str) -> Result<()> {
84        let indent = style.indent();
85
86        if text.contains('\n') {
87            let trailing_newline = text.ends_with('\n');
88            for line in text.lines() {
89                self.buffer.clear();
90                if !indent.is_empty() && !line.is_empty() {
91                    self.buffer.push_str(indent);
92                }
93                self.buffer.push_str(line);
94                self.flush(style)?;
95            }
96            if trailing_newline {
97                self.buffer.clear();
98                if !indent.is_empty() {
99                    self.buffer.push_str(indent);
100                }
101                self.flush(style)?;
102            }
103            Ok(())
104        } else {
105            self.buffer.clear();
106            if !indent.is_empty() && !text.is_empty() {
107                self.buffer.push_str(indent);
108            }
109            self.buffer.push_str(text);
110            self.flush(style)
111        }
112    }
113
114    /// Write styled text without a trailing newline
115    pub fn inline_with_style(&mut self, style: Style, text: &str) -> Result<()> {
116        if self.color {
117            write!(self.writer, "{style}{}{Reset}", text)?;
118        } else {
119            write!(self.writer, "{}", text)?;
120        }
121        self.writer.flush()?;
122        Ok(())
123    }
124
125    /// Write a line with an explicit style
126    pub fn line_with_style(&mut self, style: Style, text: &str) -> Result<()> {
127        if self.color {
128            writeln!(self.writer, "{style}{}{Reset}", text)?;
129        } else {
130            writeln!(self.writer, "{}", text)?;
131        }
132        self.writer.flush()?;
133        Ok(())
134    }
135
136    /// Write a raw line without styling
137    pub fn raw_line(&mut self, text: &str) -> Result<()> {
138        writeln!(self.writer, "{}", text)?;
139        self.writer.flush()?;
140        Ok(())
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147
148    #[test]
149    fn test_styles_construct() {
150        let info = MessageStyle::Info.style();
151        assert_eq!(info, MessageStyle::Info.style());
152        let resp = MessageStyle::Response.style();
153        assert_eq!(resp, MessageStyle::Response.style());
154        let tool = MessageStyle::Tool.style();
155        assert_eq!(tool, MessageStyle::Tool.style());
156    }
157
158    #[test]
159    fn test_renderer_buffer() {
160        let mut r = AnsiRenderer::stdout();
161        r.push("hello");
162        assert_eq!(r.buffer, "hello");
163    }
164}