Skip to main content

ansistr/
lib.rs

1use regex::Regex;
2
3/// Regex pattern for maching ANSI codes.
4pub const ANSI_REGEX: &'static str = r"[\x1b\x9b][\[()#;?]*(?:[0-9]{1,4}(?:;[0-9]{0,4})*)?[0-9A-PRZcf-nqry=><]";
5
6/// List of supported ANSI codes.
7pub const ANSI_PAIR: [[&'static str; 2]; 24] = [
8    ["\x1B[0m", "\x1B[0m"], // reset 0
9    ["\x1B[1m", "\x1B[22m"], // bold 1
10    ["\x1B[2m", "\x1B[22m"], // dim 2
11    ["\x1B[3m", "\x1B[23m"], // italic 3
12    ["\x1B[4m", "\x1B[24m"], // underline 4
13    ["\x1B[5m", "\x1B[25m"], // blink 5
14    ["\x1B[7m", "\x1B[27m"], // inverse 6
15    ["\x1B[8m", "\x1B[28m"], // hidden 7
16    ["\x1B[30m", "\x1B[39m"], // black 8
17    ["\x1B[31m", "\x1B[39m"], // red 9
18    ["\x1B[32m", "\x1B[39m"], // green 10
19    ["\x1B[33m", "\x1B[39m"], // yellow 11
20    ["\x1B[34m", "\x1B[39m"], // blue 12
21    ["\x1B[35m", "\x1B[39m"], // magenta 13
22    ["\x1B[36m", "\x1B[39m"], // cyan 14
23    ["\x1B[37m", "\x1B[39m"], // white 15
24    ["\x1B[40m", "\x1B[49m"], // bgblack 16
25    ["\x1B[41m", "\x1B[49m"], // bgred 17
26    ["\x1B[42m", "\x1B[49m"], // bggreen 18
27    ["\x1B[43m", "\x1B[49m"], // bgyellow 19
28    ["\x1B[44m", "\x1B[49m"], // bgblue 20
29    ["\x1B[45m", "\x1B[49m"], // bgmagenta 21
30    ["\x1B[46m", "\x1B[49m"], // bgcyan 22
31    ["\x1B[47m", "\x1B[49m"], // bgwhite 23
32];
33
34/// Text alignement options.
35#[derive(Debug, Clone, PartialEq)]
36pub enum TextAlign {
37    Left = 1,
38    Center = 2,
39    Right = 3,
40}
41
42/// Text alignement options.
43#[derive(Debug, Clone, PartialEq)]
44pub enum TextStyle {
45    Bold = 1,
46    Dim = 2,
47    Italic = 3,
48    Underlined = 4,
49    Blinking = 5,
50    Inversed = 6,
51    Hidden = 7,
52}
53
54/// Text alignement options.
55#[derive(Debug, Clone, PartialEq)]
56pub enum TextColor {
57    Black = 8,
58    Red = 9,
59    Green = 10,
60    Yellow = 11,
61    Blue = 12,
62    Magenta = 13,
63    Cyan = 14,
64    White = 15,
65}
66
67/// Text alignement options.
68#[derive(Debug, Clone, PartialEq)]
69pub enum TextBackground {
70    Black = 16,
71    Red = 17,
72    Green = 18,
73    Yellow = 19,
74    Blue = 20,
75    Magenta = 21,
76    Cyan = 22,
77    White = 23,
78}
79
80fn ansi_pair<'a>(code: &'a str) -> Option<&[&str; 2]> {
81    ANSI_PAIR.iter().find(|&pair| pair.iter().any(|&v| v == code))
82}
83
84/// Wraps text with ANSI style codes.
85pub fn style_str<S: Into<String>>(txt: S, style: &TextStyle) -> String {
86    let index = match style {
87        TextStyle::Bold => 1,
88        TextStyle::Dim => 2,
89        TextStyle::Italic => 3,
90        TextStyle::Underlined => 4,
91        TextStyle::Blinking => 5,
92        TextStyle::Inversed => 6,
93        TextStyle::Hidden => 7,
94    };
95    format!("{}{}{}", ANSI_PAIR[index][0], txt.into(), ANSI_PAIR[index][1])
96}
97
98/// Wraps text with ANSI color codes.
99pub fn color_str<S: Into<String>>(txt: S, color: &TextColor) -> String {
100    let index = match color {
101        TextColor::Black => 8,
102        TextColor::Red => 9,
103        TextColor::Green => 10,
104        TextColor::Yellow => 11,
105        TextColor::Blue => 12,
106        TextColor::Magenta => 13,
107        TextColor::Cyan => 14,
108        TextColor::White => 15,
109    };
110    format!("{}{}{}", ANSI_PAIR[index][0], txt.into(), ANSI_PAIR[index][1])
111}
112
113/// Wraps text with ANSI background codes.
114pub fn background_str<S: Into<String>>(txt: S, bg: &TextBackground) -> String {
115    let index = match bg {
116        TextBackground::Black => 16,
117        TextBackground::Red => 17,
118        TextBackground::Green => 18,
119        TextBackground::Yellow => 19,
120        TextBackground::Blue => 20,
121        TextBackground::Magenta => 21,
122        TextBackground::Cyan => 22,
123        TextBackground::White => 23,
124    };
125    format!("{}{}{}", ANSI_PAIR[index][0], txt.into(), ANSI_PAIR[index][1])
126}
127
128/// Strips ANSI codes from text.
129pub fn clean_str<S: Into<String>>(txt: S) -> String {
130    let txt = txt.into();
131    let regex = Regex::new(ANSI_REGEX).unwrap();
132    let clean = String::from_utf8(regex.replace_all(&txt, "").as_bytes().to_vec());
133    if clean.is_ok() {
134        clean.unwrap()
135    } else {
136        txt
137    }
138}
139
140pub fn match_indices<S: Into<String>>(txt: S) -> Vec<String> {
141    let regex = Regex::new(ANSI_REGEX).unwrap();
142    let mut result = Vec::new();
143    let mut data: String = txt.into();
144
145    loop {
146        let mat = regex.find(data.as_str());
147        if mat.is_some() {
148            let mat = mat.unwrap();
149            let start = mat.start();
150            let end = mat.end();
151            result.push(data[0..start].to_string());
152            result.push(data[start..end].to_string());
153
154            let size = data.chars().count();
155            if size == 0 {
156                break;
157            } else {
158                data = data[end..].to_string();
159            }
160        } else {
161            result.push(data);
162            break;
163        }
164    }
165
166    result
167}
168
169pub fn slice_str<S: Into<String>>(txt: S, start: usize, end: usize) -> String {
170    let mut u_start = None;
171    let mut u_end = None;
172    let mut offset = 0;
173    let mut u_offset = 0;
174    let txt = txt.into();
175
176    for chunk in match_indices(&txt).iter() {
177        let size = clean_str(chunk).len();
178        
179        if u_start.is_none() && offset + size >= start {
180            u_start = Some(u_offset + start - offset);
181        }
182        if u_end.is_none() && offset + size >= end {
183            u_end = Some(u_offset + end - offset);
184            break;
185        }
186        offset += size;
187        u_offset += chunk.len();
188    }
189
190    let u_start = match u_start {
191        Some(v) => v,
192        None => 0,
193    };
194    let u_end = match u_end {
195        Some(v) => v,
196        None => txt.len(),
197    };
198    txt[u_start..u_end].to_string()
199}
200
201pub fn size_str<S: Into<String>>(txt: S) -> usize {
202    unicode_width::UnicodeWidthStr::width(clean_str(txt).as_str())
203}
204
205pub fn pad_str<S0: Into<String>, S1: Into<String>>(txt: S0, width: usize, align: &TextAlign, chr: S1) -> String {
206    let txt = txt.into();
207    let chr = chr.into();
208
209    let size = size_str(&txt);
210    if size >= width {
211        return txt;
212    }
213
214    let chrsize = size_str(&chr);
215    let diff = width - size;
216    let (left_pad, right_pad) = match align {
217        TextAlign::Left => (0, diff / chrsize),
218        TextAlign::Right => (diff / chrsize, 0),
219        TextAlign::Center => (diff / chrsize / 2, diff - diff / chrsize / 2),
220    };
221
222    let mut result = String::new();
223    for _ in 0..left_pad {
224        result.push_str(&chr);
225    }
226    result.push_str(&txt);
227    for _ in 0..right_pad {
228        result.push_str(&chr);
229    }
230    result
231}
232
233pub fn trucate_str<S0: Into<String>, S1: Into<String>>(txt: S0, width: usize, align: &TextAlign, tail: S1) -> String {
234    let txt = txt.into();
235    let tail = tail.into();
236
237    let size = size_str(&txt);
238    if width >= size {
239        return txt;
240    }
241
242    let t_size = size_str(&tail);
243    match align {
244        TextAlign::Left => {
245            let text = slice_str(&txt, 0, width - t_size).trim().to_string();
246            format!("{}{}", text, tail)
247        },
248        TextAlign::Right => {
249            let text = slice_str(&txt, size - width + t_size, size).trim().to_string();
250            format!("{}{}", tail, text)
251        },
252        TextAlign::Center => {
253            let dim = (width - t_size) / 2;
254            let left = slice_str(&txt, 0, dim).trim().to_string();
255            let right = slice_str(&txt, size - width + t_size + dim, size).trim().to_string();
256            format!("{}{}{}", left, tail, right)
257        },
258    }
259}
260
261pub fn wrap_str<S: Into<String>>(txt: S, width: usize) -> String {
262    let mut result: Vec<String> = Vec::new();
263    let txt = txt.into();
264
265    for line in txt.lines() {
266        let mut words: Vec<String> = Vec::new();
267        let mut length = 0;
268
269        for (wcount, word) in line.split(" ").enumerate() {
270            let word_size = size_str(word);
271            if length + word_size >= width && words.len() > 0 {
272                result.push(words.join(" "));
273                words =  Vec::new();
274                length = 0;
275            }
276            length += word_size + if wcount > 0 { 1 } else { 0 }; // include spaces
277            words.push(word.to_string());
278        }
279
280        if words.len() > 0 {
281            result.push(words.join(" "));
282        }
283    }
284
285    result.join("\n")
286}
287
288pub fn repaire_str<S: Into<String>>(txt: S) -> String {
289    let mut ansis: Vec<Vec<String>> = Vec::new();
290    let txt = txt.into();
291
292    let lines: Vec<String> = txt.split("\n").map(|line| {
293        let parts = match_indices(line);
294
295        let mut result: Vec<String> = Vec::new();
296        let ansiiter = &ansis;
297        for ansi in ansiiter.into_iter() {
298            result.push(ansi[0].to_string());
299        }
300        for part in parts.into_iter() {
301            let pair = ansi_pair(part.as_str());
302            if pair.is_some() {
303                let pair = pair.unwrap();
304                let opentag = pair[0].to_string();
305                let closetag = pair[1].to_string();
306                if part == opentag {
307                    ansis.push(vec![opentag, closetag]);
308                } else if part == closetag {
309                    let index = ansis.iter().position(|a| a[1].to_string() == closetag);
310                    if index.is_some() {
311                        ansis.remove(index.unwrap());
312                    }
313                }
314            }
315            result.push(part.to_string());
316        }
317        let ansiiter = &ansis;
318        for ansi in ansiiter.into_iter() {
319            result.push(ansi[1].to_string());
320        }
321        result.join("")
322    }).collect();
323
324    lines.join("\n")
325}
326
327/// Unit tests.
328#[cfg(test)]
329mod tests {
330    use super::*;
331
332    #[test]
333    fn finds_ansi_pair() {
334        assert_eq!(ansi_pair(&ANSI_PAIR[0][1]), Some(&ANSI_PAIR[0]));
335        assert_eq!(ansi_pair("foo"), None);
336    }
337
338    #[test]
339    fn applies_ansi_style() {
340        style_str("foo", &TextStyle::Bold);
341        assert_eq!(
342            style_str("foo", &TextStyle::Bold),
343            format!("{}{}{}", "\x1B[1m", "foo", "\x1B[22m"),
344        );
345    }
346
347    #[test]
348    fn applies_ansi_color() {
349        assert_eq!(
350            color_str("foo", &TextColor::Red),
351            format!("{}{}{}", "\x1B[31m", "foo", "\x1B[39m"),
352        );
353    }
354
355    #[test]
356    fn applies_ansi_background() {
357        assert_eq!(
358            background_str("foo", &TextBackground::Green),
359            format!("{}{}{}", "\x1B[42m", "foo", "\x1B[49m"),
360        );
361    }
362
363    #[test]
364    fn strips_ansi_codes() {
365        assert_eq!(clean_str("aaa\x1B[0mbbb\x1B[0mccc"), "aaabbbccc");
366    }
367
368    #[test]
369    fn matches_ansi_indices() {
370        assert_eq!(match_indices("This is\x1B[39m long"), vec!["This is", "\x1B[39m", " long"]);
371        assert_eq!(match_indices("This is\x1B[39m long \x1B[46mtext for test"), vec!["This is", "\x1B[39m", " long ", "\x1B[46m", "text for test"]);
372    }
373
374    #[test]
375    fn slices_ansi_str() {
376        assert_eq!(slice_str("a\x1B[32maa\x1B[32mb\x1B[32mbb\x1B[32mcccdddeeefff", 5, 10), "b\x1B[32mcccd");
377    }
378
379    #[test]
380    fn sizes_ansi_str() {
381        assert_eq!(size_str("aaa\x1B[0mbbb\x1B[0mccc"), 9);
382    }
383
384    #[test]
385    fn pads_ansi_str() {
386        assert_eq!(pad_str("fo\x1B[39mobar", 10, &TextAlign::Left, "+"), "fo\x1B[39mobar++++");
387        assert_eq!(pad_str("fo\x1B[39mobar", 10, &TextAlign::Right, "+"), "++++fo\x1B[39mobar");
388        assert_eq!(pad_str("fo\x1B[39mobar", 10, &TextAlign::Center, "+"), "++fo\x1B[39mobar++");
389        assert_eq!(pad_str("fo\x1B[39mobar", 10, &TextAlign::Left, "\x1B[39m+!"), "fo\x1B[39mobar\x1B[39m+!\x1B[39m+!");
390    }
391
392    #[test]
393    fn truncates_ansi_str() {
394        assert_eq!(trucate_str("fo\x1B[39mobarbaz", 5, &TextAlign::Left, "+"), "fo\x1B[39mob+");
395        assert_eq!(trucate_str("fo\x1B[39mobarbaz", 5, &TextAlign::Right, "+++"), "+++az");
396        assert_eq!(trucate_str("fo\x1B[39mobarbaz", 5, &TextAlign::Center, "+++"), "f+++z");
397    }
398
399    #[test]
400    fn wraps_ansi_str() {
401        assert_eq!(wrap_str("This is \x1B[39ma very long tekst for testing\x1B[39m only.", 10), vec![
402            "This is \x1B[39ma",
403            "very long",
404            "tekst for",
405            "testing\x1B[39m",
406            "only."
407        ].join("\n"));
408    }
409
410    #[test]
411    fn repairs_multiline_ansi_str() {
412        assert_eq!(repaire_str(&vec![
413            "This is \x1B[31mlong",
414            "string 利干 sample",
415            "this is 利干 sample\x1B[39m long code",
416        ].join("\n")), vec![
417            "This is \x1B[31mlong\x1B[39m",
418            "\x1B[31mstring 利干 sample\x1B[39m",
419            "\x1B[31mthis is 利干 sample\x1B[39m long code",
420        ].join("\n"));
421    }
422}