deepseek_rust_cli/tui/colorizer/
mod.rs

1//! Streaming colorizer for reasoning/content output.
2
3pub mod highlighter;
4pub mod types;
5pub mod utils;
6
7use std::fmt::Write as FmtWrite;
8
9pub use highlighter::CodeColorizer;
10pub use types::CodeLang;
11pub use utils::truncate_result;
12
13use crate::tui::colorizer::types::State;
14
15pub struct StreamColorizer {
16    state: State,
17    /// Pending output that might be part of an unclosed construct
18    pending: String,
19    dimmed: bool,
20    first_feed: bool,
21}
22
23impl Default for StreamColorizer {
24    fn default() -> Self {
25        Self::new()
26    }
27}
28
29impl StreamColorizer {
30    pub fn new() -> Self {
31        Self {
32            state: State::Normal,
33            pending: String::new(),
34            dimmed: false,
35            first_feed: true,
36        }
37    }
38
39    pub fn set_dimmed(&mut self, dimmed: bool) {
40        self.dimmed = dimmed;
41    }
42
43    fn reset_code(&self) -> String {
44        if self.dimmed {
45            "\x1b[0m\x1b[2m".to_string()
46        } else {
47            "\x1b[0m".to_string()
48        }
49    }
50
51    /// Feed a chunk of text, return colored output ready to print.
52    /// Call `finish()` at the end to flush remaining buffer.
53    pub fn feed(&mut self, chunk: &str) -> String {
54        let input = format!("{}{}", self.pending, chunk);
55        self.pending.clear();
56
57        let mut out = String::new();
58        if self.first_feed && self.dimmed {
59            out.push_str("\x1b[2m");
60            self.first_feed = false;
61        }
62        let chars: Vec<char> = input.chars().collect();
63        let len = chars.len();
64        let mut i = 0;
65
66        while i < len {
67            match self.state {
68                State::Normal => {
69                    // Look for backtick or file path
70                    if chars[i] == '`' {
71                        // Check for fenced block (```)
72                        if i + 2 < len && chars[i + 1] == '`' && chars[i + 2] == '`' {
73                            // Opening fenced block
74                            let _ = write!(out, "\x1b[33m```{}", self.reset_code()); // yellow backticks
75                            i += 3;
76                            // Read language tag
77                            let mut lang = String::new();
78                            while i < len && chars[i] != '\n' && chars[i] != '\r' {
79                                lang.push(chars[i]);
80                                i += 1;
81                            }
82                            if !lang.is_empty() {
83                                let _ = write!(out, "\x1b[36m{}{}", lang, self.reset_code());
84                                // cyan lang
85                            }
86                            // Skip newline after lang tag
87                            if i < len && chars[i] == '\r' {
88                                i += 1;
89                            }
90                            if i < len && chars[i] == '\n' {
91                                out.push('\n');
92                                i += 1;
93                            }
94                            self.state = State::FencedBlock {
95                                lang: lang.trim().to_string(),
96                            };
97                        } else {
98                            // Start inline code
99                            let _ = write!(out, "\x1b[32m`{}", self.reset_code()); // green backtick
100                            i += 1;
101                            self.state = State::InlineCode;
102                        }
103                    } else if self.is_path_boundary(&chars, i, len) {
104                        // Try to match a file path
105                        let path_end = self.match_path(&chars, i, len);
106                        if path_end > i {
107                            let path: String = chars[i..path_end].iter().collect();
108                            let _ = write!(out, "\x1b[34m{}{}", path, self.reset_code()); // blue file path
109                            i = path_end;
110                        } else {
111                            out.push(chars[i]);
112                            i += 1;
113                        }
114                    } else {
115                        out.push(chars[i]);
116                        i += 1;
117                    }
118                }
119                State::InlineCode => {
120                    if chars[i] == '`' {
121                        let _ = write!(out, "\x1b[32m`{}", self.reset_code()); // closing green backtick
122                        i += 1;
123                        self.state = State::Normal;
124                    } else {
125                        // Content inside inline code - keep green
126                        out.push_str("\x1b[32m");
127                        while i < len && chars[i] != '`' {
128                            out.push(chars[i]);
129                            i += 1;
130                        }
131                        out.push_str(&self.reset_code());
132                    }
133                }
134                State::FencedBlock { ref lang } => {
135                    // Look for closing ```
136                    if chars[i] == '`' && i + 2 < len && chars[i + 1] == '`' && chars[i + 2] == '`'
137                    {
138                        let _ = write!(out, "\x1b[33m```{}", self.reset_code()); // yellow closing
139                        i += 3;
140                        self.state = State::Normal;
141                    } else {
142                        // Content inside code block - dim white
143                        let lang_clone = lang.clone();
144                        out.push_str("\x1b[37m");
145                        while i < len {
146                            if chars[i] == '`'
147                                && i + 2 < len
148                                && chars[i + 1] == '`'
149                                && chars[i + 2] == '`'
150                            {
151                                break;
152                            }
153                            out.push(chars[i]);
154                            i += 1;
155                        }
156                        out.push_str(&self.reset_code());
157                        self.state = State::FencedBlock { lang: lang_clone };
158                    }
159                }
160            }
161        }
162
163        out
164    }
165
166    /// Call when streaming is done to flush any remaining state
167    pub fn finish(&mut self) -> String {
168        let mut out = String::new();
169
170        // Close any open constructs
171        match self.state {
172            State::InlineCode => {
173                let _ = write!(out, "\x1b[32m`{}", self.reset_code()); // close inline
174            }
175            State::FencedBlock { .. } => {
176                let _ = write!(out, "\x1b[33m```{}", self.reset_code()); // close block
177            }
178            _ => {}
179        }
180
181        if !self.pending.is_empty() {
182            out.push_str(&self.pending);
183            self.pending.clear();
184        }
185
186        if self.dimmed {
187            out.push_str("\x1b[0m");
188        }
189
190        self.state = State::Normal;
191        self.first_feed = true;
192        out
193    }
194
195    /// Check if we're at a potential file path boundary
196    fn is_path_boundary(&self, chars: &[char], i: usize, _len: usize) -> bool {
197        let c = chars[i];
198        // Path must start with ./ or ../ or / or a word char that looks like a file
199        if c == '.' || c == '/' || c == '~' {
200            return true;
201        }
202        if c.is_alphanumeric() {
203            // Check if previous char is whitespace or boundary
204            if i == 0 || chars[i - 1].is_whitespace() || chars[i - 1] == '(' || chars[i - 1] == '['
205            {
206                // Look ahead for path-like patterns
207                return true;
208            }
209        }
210        false
211    }
212
213    /// Try to match a file path starting at position i
214    fn match_path(&self, chars: &[char], start: usize, len: usize) -> usize {
215        let mut end = start;
216
217        // Common path prefixes
218        if start < len {
219            match chars[start] {
220                '/' => {
221                    end += 1;
222                }
223                '~' if start + 1 < len && chars[start + 1] == '/' => {
224                    end += 2;
225                }
226                '.' if start + 1 < len && chars[start + 1] == '/' => {
227                    end += 2;
228                }
229                c if c.is_alphanumeric() => {
230                    // Must contain a path separator or known extension
231                }
232                _ => return start,
233            }
234        }
235
236        // Continue matching path characters
237        while end < len {
238            let c = chars[end];
239            if c.is_alphanumeric()
240                || c == '/'
241                || c == '.'
242                || c == '-'
243                || c == '_'
244                || c == ' '
245                || c == '~'
246                || c == '+'
247                || c == '@'
248                || c == '#'
249                || c == ':'
250            {
251                // Check for common file extensions to stop at
252                if c == '.' && end + 1 < len {
253                    // Look for known extensions
254                    let remaining = &chars[end + 1..];
255                    let remaining_str: String = remaining.iter().collect();
256                    let ext_candidates = [
257                        "rs",
258                        "py",
259                        "js",
260                        "ts",
261                        "go",
262                        "java",
263                        "c",
264                        "cpp",
265                        "h",
266                        "hpp",
267                        "rb",
268                        "php",
269                        "swift",
270                        "kt",
271                        "scala",
272                        "sh",
273                        "bash",
274                        "zsh",
275                        "fish",
276                        "ps1",
277                        "toml",
278                        "yaml",
279                        "yml",
280                        "json",
281                        "xml",
282                        "html",
283                        "css",
284                        "scss",
285                        "md",
286                        "txt",
287                        "log",
288                        "csv",
289                        "env",
290                        "cfg",
291                        "conf",
292                        "lock",
293                        "gitignore",
294                        "dockerfile",
295                        "nix",
296                        "lua",
297                        "vim",
298                        "el",
299                        "ex",
300                        "exs",
301                        "erl",
302                        "hrl",
303                        "sql",
304                        "graphql",
305                        "proto",
306                        "vue",
307                        "svelte",
308                        "tsx",
309                        "jsx",
310                        "mjs",
311                        "wasm",
312                        "wat",
313                        "bc",
314                        "dc",
315                        "awk",
316                        "sed",
317                    ];
318                    let matched = ext_candidates.iter().any(|ext| {
319                        remaining_str.len() >= ext.len()
320                            && remaining_str[..ext.len()].eq_ignore_ascii_case(ext)
321                            && (remaining_str.len() == ext.len()
322                                || remaining_str
323                                    .as_bytes()
324                                    .get(ext.len())
325                                    .is_none_or(|&b| !b.is_ascii_alphanumeric() && b != b'_'))
326                    });
327                    if matched {
328                        // Include the dot
329                        end += 1;
330                        // Include the extension
331                        while end < len && chars[end].is_alphanumeric() {
332                            end += 1;
333                        }
334                        // Also include trailing / if present
335                        if end < len && chars[end] == '/' {
336                            end += 1;
337                            continue;
338                        }
339                        break;
340                    }
341                }
342
343                // Stop at whitespace if we have a valid path
344                if c.is_whitespace() {
345                    break;
346                }
347
348                // Stop at certain punctuation
349                if c == ','
350                    || c == ';'
351                    || c == ')'
352                    || c == ']'
353                    || c == '}'
354                    || c == '"'
355                    || c == '\''
356                    || c == '>'
357                    || c == '`'
358                {
359                    break;
360                }
361
362                end += 1;
363            } else {
364                break;
365            }
366        }
367
368        // Must have at least 3 chars and look like a path
369        if end - start >= 3 {
370            let segment: String = chars[start..end].iter().collect();
371            if segment.contains('/')
372                || segment.contains('.')
373                || segment.ends_with("rc")
374                || segment.ends_with("file")
375            {
376                return end;
377            }
378        }
379
380        start // no match
381    }
382}
383
384#[cfg(test)]
385mod tests {
386    use super::*;
387
388    #[test]
389    fn test_inline_code() {
390        let mut c = StreamColorizer::new();
391        let out = c.feed("Use `cargo build` to compile.");
392        assert!(out.contains("\x1b[32m"));
393    }
394
395    #[test]
396    fn test_file_path() {
397        let mut c = StreamColorizer::new();
398        let out = c.feed("Edit src/main.rs and Cargo.toml");
399        assert!(out.contains("\x1b[34m"));
400    }
401
402    #[test]
403    fn test_fenced_block() {
404        let mut c = StreamColorizer::new();
405        let out = c.feed("```rust\nlet x = 1;\n```");
406        assert!(out.contains("\x1b[33m"));
407        assert!(out.contains("\x1b[36mrust\x1b[0m"));
408    }
409
410    #[test]
411    fn test_code_rust_keywords() {
412        let code = "fn main() {\n    let x = 42;\n}";
413        let colored = CodeColorizer::highlight(code, CodeLang::Rust, None);
414        assert!(colored.contains("\x1b[34mfn\x1b[0m"));
415        assert!(colored.contains("\x1b[34mlet\x1b[0m"));
416        assert!(colored.contains("\x1b[35m42\x1b[0m"));
417    }
418
419    #[test]
420    fn test_lang_from_path() {
421        assert_eq!(CodeLang::from_path("src/main.rs"), CodeLang::Rust);
422        assert_eq!(CodeLang::from_path("app.py"), CodeLang::Python);
423        assert_eq!(CodeLang::from_path("unknown.xyz"), CodeLang::Generic);
424    }
425
426    #[test]
427    fn test_truncate_result() {
428        let short = "hello";
429        assert_eq!(truncate_result(short, 100), "hello");
430
431        let long = "a".repeat(200);
432        let truncated = truncate_result(&long, 50);
433        assert!(truncated.len() <= 120);
434    }
435}
deepseek_rust_cli/tui/colorizer/mod.rs

deepseek_rust_cli/tui/colorizer/
mod.rs
