sql_cli/utils/
string_utils.rs

1use unicode_width::UnicodeWidthStr;
2
3/// Strip ANSI escape codes from a string and return the display width
4/// This handles ANSI SGR (Select Graphic Rendition) codes like colors and styles
5/// and correctly calculates Unicode character widths (e.g., emoji take 2 columns)
6///
7/// # Examples
8///
9/// ```
10/// use sql_cli::utils::string_utils::display_width;
11///
12/// // Simple ASCII text
13/// assert_eq!(display_width("hello"), 5);
14///
15/// // ANSI colored text (escape codes don't count toward width)
16/// assert_eq!(display_width("\x1b[31mred\x1b[0m"), 3);
17///
18/// // Unicode emoji (takes 2 columns)
19/// assert_eq!(display_width("→"), 1);  // Right arrow
20/// assert_eq!(display_width("⚡"), 2);  // Lightning bolt emoji
21/// ```
22pub fn display_width(s: &str) -> usize {
23    let mut result = String::new();
24    let mut chars = s.chars().peekable();
25
26    while let Some(ch) = chars.next() {
27        if ch == '\x1b' {
28            // Check for ANSI escape sequence
29            if chars.peek() == Some(&'[') {
30                chars.next(); // consume '['
31                              // Skip until we find a letter (the command character)
32                while let Some(&next_ch) = chars.peek() {
33                    chars.next();
34                    if next_ch.is_ascii_alphabetic() {
35                        break;
36                    }
37                }
38            } else {
39                result.push(ch);
40            }
41        } else {
42            result.push(ch);
43        }
44    }
45
46    // Use unicode-width to get the actual display width
47    // This correctly handles:
48    // - ASCII characters (width 1)
49    // - Emoji and other wide characters (width 2)
50    // - Zero-width characters (width 0)
51    // - Combining characters
52    result.width()
53}
54
55#[cfg(test)]
56mod tests {
57    use super::*;
58
59    #[test]
60    fn test_display_width_ascii() {
61        assert_eq!(display_width("hello"), 5);
62        assert_eq!(display_width("test"), 4);
63        assert_eq!(display_width(""), 0);
64    }
65
66    #[test]
67    fn test_display_width_ansi_codes() {
68        // Red colored text
69        assert_eq!(display_width("\x1b[31mred\x1b[0m"), 3);
70        // Bold text
71        assert_eq!(display_width("\x1b[1mbold\x1b[0m"), 4);
72        // RGB colored text
73        assert_eq!(display_width("\x1b[38;2;255;0;0mRGB red\x1b[0m"), 7);
74    }
75
76    #[test]
77    fn test_display_width_unicode() {
78        // Single-width Unicode
79        assert_eq!(display_width("→"), 1);
80        assert_eq!(display_width("café"), 4);
81
82        // Double-width characters (emoji)
83        assert_eq!(display_width("⚡"), 2);
84        assert_eq!(display_width("😀"), 2);
85    }
86
87    #[test]
88    fn test_display_width_mixed() {
89        // ANSI + Unicode
90        assert_eq!(
91            display_width("\x1b[32m→ CONSEC\x1b[0m"),
92            8 // 1 (arrow) + 1 (space) + 6 (CONSEC)
93        );
94        assert_eq!(
95            display_width("\x1b[31m⚡ REPEAT\x1b[0m"),
96            9 // 2 (bolt) + 1 (space) + 6 (REPEAT)
97        );
98    }
99}
sql_cli/utils/string_utils.rs

sql_cli/utils/
string_utils.rs
