vtcode_core/utils/
ansi.rs1use 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#[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
40pub struct AnsiRenderer {
42 writer: AutoStream<io::Stdout>,
43 buffer: String,
44 color: bool,
45}
46
47impl AnsiRenderer {
48 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 pub fn push(&mut self, text: &str) {
66 self.buffer.push_str(text);
67 }
68
69 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 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 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 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 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}