1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

use std::borrow::Cow;

use regex::{Matches, Regex};

lazy_static::lazy_static! {
    static ref STRIP_ANSI_RE: Regex =
        Regex::new(r"[\x1b\x9b][\[()#;?]*(?:[0-9]{1,4}(?:;[0-9]{0,4})*)?[0-9A-PRZcf-nqry=><]")
            .unwrap();
}

/// Helper function to strip ansi codes.
pub fn strip_ansi_codes(s: &str) -> Cow<str> {
    STRIP_ANSI_RE.replace_all(s, "")
}

/// An iterator over ansi codes in a string.
///
/// This type can be used to scan over ansi codes in a string.
/// It yields tuples in the form `(s, is_ansi)` where `s` is a slice of
/// the original string and `is_ansi` indicates if the slice contains
/// ansi codes or string values.
pub struct AnsiCodeIterator<'a> {
    s: &'a str,
    pending_item: Option<(&'a str, bool)>,
    last_idx: usize,
    cur_idx: usize,
    iter: Matches<'static, 'a>,
}

impl<'a> AnsiCodeIterator<'a> {
    /// Creates a new ansi code iterator.
    pub fn new(s: &'a str) -> AnsiCodeIterator<'a> {
        AnsiCodeIterator {
            s,
            pending_item: None,
            last_idx: 0,
            cur_idx: 0,
            iter: STRIP_ANSI_RE.find_iter(s),
        }
    }

    /// Returns the string slice up to the current match.
    pub fn current_slice(&self) -> &str {
        &self.s[..self.cur_idx]
    }

    /// Returns the string slice from the current match to the end.
    pub fn rest_slice(&self) -> &str {
        &self.s[self.cur_idx..]
    }
}

impl<'a> Iterator for AnsiCodeIterator<'a> {
    type Item = (&'a str, bool);

    fn next(&mut self) -> Option<(&'a str, bool)> {
        if let Some(pending_item) = self.pending_item.take() {
            self.cur_idx += pending_item.0.len();
            Some(pending_item)
        } else if let Some(m) = self.iter.next() {
            let s = &self.s[self.last_idx..m.start()];
            self.last_idx = m.end();
            if s.is_empty() {
                self.cur_idx = m.end();
                Some((m.as_str(), true))
            } else {
                self.cur_idx = m.start();
                self.pending_item = Some((m.as_str(), true));
                Some((s, false))
            }
        } else if self.last_idx < self.s.len() {
            let rv = &self.s[self.last_idx..];
            self.cur_idx = self.s.len();
            self.last_idx = self.s.len();
            Some((rv, false))
        } else {
            None
        }
    }
}

#[test]
fn test_ansi_iter_re() {
    use crate::style;
    let s = format!("Hello {}!", style("World").red().force_styling(true));
    let mut iter = AnsiCodeIterator::new(&s);
    assert_eq!(iter.next(), Some(("Hello ", false)));
    assert_eq!(iter.current_slice(), "Hello ");
    assert_eq!(iter.rest_slice(), "\x1b[31mWorld\x1b[0m!");
    assert_eq!(iter.next(), Some(("\x1b[31m", true)));
    assert_eq!(iter.current_slice(), "Hello \x1b[31m");
    assert_eq!(iter.rest_slice(), "World\x1b[0m!");
    assert_eq!(iter.next(), Some(("World", false)));
    assert_eq!(iter.current_slice(), "Hello \x1b[31mWorld");
    assert_eq!(iter.rest_slice(), "\x1b[0m!");
    assert_eq!(iter.next(), Some(("\x1b[0m", true)));
    assert_eq!(iter.current_slice(), "Hello \x1b[31mWorld\x1b[0m");
    assert_eq!(iter.rest_slice(), "!");
    assert_eq!(iter.next(), Some(("!", false)));
    assert_eq!(iter.current_slice(), "Hello \x1b[31mWorld\x1b[0m!");
    assert_eq!(iter.rest_slice(), "");
    assert_eq!(iter.next(), None);
}

#[test]
fn test_ansi_iter_re_on_multi() {
    use crate::style;
    let s = format!("{}", style("a").red().bold().force_styling(true));
    let mut iter = AnsiCodeIterator::new(&s);
    assert_eq!(iter.next(), Some(("\x1b[31m", true)));
    assert_eq!(iter.current_slice(), "\x1b[31m");
    assert_eq!(iter.rest_slice(), "\x1b[1ma\x1b[0m");
    assert_eq!(iter.next(), Some(("\x1b[1m", true)));
    assert_eq!(iter.current_slice(), "\x1b[31m\x1b[1m");
    assert_eq!(iter.rest_slice(), "a\x1b[0m");
    assert_eq!(iter.next(), Some(("a", false)));
    assert_eq!(iter.current_slice(), "\x1b[31m\x1b[1ma");
    assert_eq!(iter.rest_slice(), "\x1b[0m");
    assert_eq!(iter.next(), Some(("\x1b[0m", true)));
    assert_eq!(iter.current_slice(), "\x1b[31m\x1b[1ma\x1b[0m");
    assert_eq!(iter.rest_slice(), "");
    assert_eq!(iter.next(), None);
}