Skip to main content

nils_common/
shell.rs

1use std::borrow::Cow;
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4pub enum AnsiStripMode {
5    CsiSgrOnly,
6    CsiAnyTerminator,
7}
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum SingleQuoteEscapeStyle {
11    Backslash,
12    DoubleQuoteBoundary,
13}
14
15pub fn quote_posix_single(input: &str) -> String {
16    quote_posix_single_with_style(input, SingleQuoteEscapeStyle::Backslash)
17}
18
19pub fn quote_posix_single_with_style(input: &str, style: SingleQuoteEscapeStyle) -> String {
20    if input.is_empty() {
21        return "''".to_string();
22    }
23
24    let mut out = String::from("'");
25    for ch in input.chars() {
26        if ch == '\'' {
27            match style {
28                SingleQuoteEscapeStyle::Backslash => out.push_str("'\\''"),
29                SingleQuoteEscapeStyle::DoubleQuoteBoundary => out.push_str("'\"'\"'"),
30            }
31        } else {
32            out.push(ch);
33        }
34    }
35    out.push('\'');
36    out
37}
38
39pub fn strip_ansi(input: &str, mode: AnsiStripMode) -> Cow<'_, str> {
40    let bytes = input.as_bytes();
41    let mut i = 0usize;
42    let mut copied_from = 0usize;
43    let mut out: Option<String> = None;
44
45    while i < bytes.len() {
46        if bytes[i] == 0x1b && i + 1 < bytes.len() && bytes[i + 1] == b'[' {
47            let mut j = i + 2;
48            let mut should_strip = false;
49            match mode {
50                AnsiStripMode::CsiSgrOnly => {
51                    while j < bytes.len() {
52                        let b = bytes[j];
53                        j += 1;
54                        if (0x40..=0x7e).contains(&b) {
55                            should_strip = b == b'm';
56                            break;
57                        }
58                    }
59                }
60                AnsiStripMode::CsiAnyTerminator => {
61                    while j < bytes.len() {
62                        let b = bytes[j];
63                        j += 1;
64                        if (0x40..=0x7e).contains(&b) {
65                            should_strip = true;
66                            break;
67                        }
68                    }
69                }
70            }
71
72            if should_strip {
73                let buffer = out.get_or_insert_with(|| String::with_capacity(input.len()));
74                buffer.push_str(&input[copied_from..i]);
75                copied_from = j;
76                i = j;
77                continue;
78            }
79        }
80
81        i += 1;
82    }
83
84    if let Some(mut buffer) = out {
85        buffer.push_str(&input[copied_from..]);
86        Cow::Owned(buffer)
87    } else {
88        Cow::Borrowed(input)
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::SingleQuoteEscapeStyle;
95    use super::{AnsiStripMode, quote_posix_single, quote_posix_single_with_style, strip_ansi};
96    use std::borrow::Cow;
97
98    #[test]
99    fn quote_posix_single_uses_backslash_style() {
100        assert_eq!(quote_posix_single("a'b"), "'a'\\''b'");
101    }
102
103    #[test]
104    fn quote_posix_single_with_double_quote_boundary_style() {
105        let out = quote_posix_single_with_style("a'b", SingleQuoteEscapeStyle::DoubleQuoteBoundary);
106        assert_eq!(out, "'a'\"'\"'b'");
107    }
108
109    #[test]
110    fn quote_posix_single_handles_empty_input() {
111        assert_eq!(quote_posix_single(""), "''");
112    }
113
114    #[test]
115    fn strip_ansi_sgr_removes_m_sequences() {
116        let input = "\x1b[31mred\x1b[0m plain";
117        assert_eq!(strip_ansi(input, AnsiStripMode::CsiSgrOnly), "red plain");
118    }
119
120    #[test]
121    fn strip_ansi_any_terminator_removes_k_sequence() {
122        let input = "a\x1b[2Kb";
123        assert_eq!(strip_ansi(input, AnsiStripMode::CsiAnyTerminator), "ab");
124    }
125
126    #[test]
127    fn strip_ansi_sgr_only_keeps_non_sgr_csi_sequences() {
128        let input = "a\x1b[2Kb";
129        assert_eq!(strip_ansi(input, AnsiStripMode::CsiSgrOnly), input);
130    }
131
132    #[test]
133    fn strip_ansi_sgr_only_keeps_incomplete_csi_sequences() {
134        let input = "a\x1b[31";
135        assert_eq!(strip_ansi(input, AnsiStripMode::CsiSgrOnly), input);
136    }
137
138    #[test]
139    fn strip_ansi_returns_borrowed_when_no_escape_found() {
140        let input = "plain text";
141        let out = strip_ansi(input, AnsiStripMode::CsiSgrOnly);
142        assert!(matches!(out, Cow::Borrowed("plain text")));
143    }
144}