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