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}