rusty_prompt/
writer.rs

1use std::collections::HashMap;
2use std::io;
3use std::io::{stdout, Write};
4
5use lazy_static::lazy_static;
6use crate::colored_string::{ColoredChar, ColoredStr, ColoredString};
7
8#[derive(Copy, Clone, Hash, Ord, PartialOrd, Eq, PartialEq)]
9pub enum DisplayAttribute {
10    Reset,
11    Bold,
12    LowIntensity,
13    Italic,
14    Underline,
15    Blink,
16    RapidBlink,
17    Reverse,
18    Invisible,
19    CrossedOut,
20    DefaultFont,
21}
22
23impl Default for DisplayAttribute {
24    fn default() -> Self {
25        Self::Reset
26    }
27}
28
29#[derive(Copy, Clone, Hash, Ord, PartialOrd, Eq, PartialEq)]
30pub enum Color {
31    Default,
32    Black,
33    DarkRed,
34    DarkGreen,
35    Brown,
36    DarkBlue,
37    Purple,
38    Cyan,
39    LightGray,
40    DarkGray,
41    Red,
42    Green,
43    Yellow,
44    Blue,
45    Fuchsia,
46    Turquoise,
47    White,
48}
49
50impl From<Option<Color>> for Color {
51    fn from(c: Option<Color>) -> Self {
52        match c {
53            None => Color::Default,
54            Some(c) => c,
55        }
56    }
57}
58
59impl Default for Color {
60    fn default() -> Self {
61        Self::Default
62    }
63}
64
65pub struct Writer {
66    buffer: Vec<u8>,
67    changed: bool,
68}
69
70impl Writer {
71    pub(crate) fn new() -> Self {
72        let mut writer = Self {
73            buffer: vec![],
74            changed: true,
75        };
76        writer.write_raw(&['\n' as u8, '\r' as u8]);
77        writer
78    }
79}
80
81impl Writer {
82    pub(crate) fn flush(&mut self) -> io::Result<()> {
83        let mut out = stdout().lock();
84        out.write_all(&self.buffer)?;
85        out.flush().unwrap();
86        self.buffer.clear();
87        self.changed = true;
88        Ok(())
89    }
90}
91
92impl Writer {
93    pub fn write_raw(&mut self, data: &[u8]) {
94        self.buffer.extend_from_slice(data);
95        self.changed = true;
96    }
97
98    pub fn write(&mut self, data: &[u8]) {
99        let replaced = replace_vec(data, 0x1b, Some('?' as u8));
100        self.write_raw(&replaced);
101    }
102
103    pub fn write_raw_str(&mut self, data: &str) {
104        self.write_raw(data.as_bytes());
105    }
106
107    pub fn write_str(&mut self, data: &str) {
108        self.write(data.as_bytes());
109    }
110
111    pub fn write_colored_string(&mut self, data: &ColoredString) {
112        self.write_colored_chars(data.get_all_chars());
113    }
114
115    pub fn write_colored_chars(&mut self, data: &[ColoredChar]) {
116        self.set_display_attributes(Color::Default, Color::Default, DisplayAttribute::Reset);
117        let mut prev_attributes = (Color::Default, Color::Default, DisplayAttribute::Reset);
118        for chr in data {
119            if chr.get_foreground_color() != prev_attributes.0 ||
120                chr.get_background_color() != prev_attributes.1 ||
121                chr.get_attribute() != prev_attributes.2 {
122                prev_attributes = (chr.get_foreground_color(), chr.get_background_color(), chr.get_attribute());
123                self.set_display_attributes(chr.get_foreground_color(), chr.get_background_color(), chr.get_attribute());
124            }
125            self.write(&[chr.get_char()]);
126        }
127        self.set_display_attributes(Color::Default, Color::Default, DisplayAttribute::Reset);
128    }
129
130    pub fn write_colored_str(&mut self, data: &ColoredStr) {
131        self.write_colored_chars(data.get_all_chars());
132    }
133}
134
135impl Writer {
136    pub fn erase_display(&mut self) {
137        self.write_raw(&[0x1b, '[' as u8, '2' as u8, 'J' as u8]);
138    }
139
140    pub fn erase_up(&mut self) {
141        self.write_raw(&[0x1b, '[' as u8, '1' as u8, 'J' as u8]);
142    }
143
144    pub fn erase_down(&mut self) {
145        self.write_raw(&[0x1b, '[' as u8, 'J' as u8]);
146    }
147
148    pub fn erase_start_of_line(&mut self) {
149        self.write_raw(&[0x1b, '[' as u8, '1' as u8, 'K' as u8]);
150    }
151
152    pub fn erase_end_of_line(&mut self) {
153        self.write_raw(&[0x1b, '[' as u8, 'K' as u8]);
154    }
155
156    pub fn erase_line(&mut self) {
157        self.write_raw(&[0x1b, '[' as u8, '2' as u8, 'K' as u8]);
158    }
159}
160
161impl Writer {
162    pub fn show_cursor(&mut self) {
163        self.write_raw(&[
164            0x1b, '[' as u8, '?' as u8, '1' as u8, '2' as u8, 'l' as u8, 0x1b, '[' as u8,
165            '?' as u8, '2' as u8, '5' as u8, 'h' as u8,
166        ]);
167    }
168
169    pub fn hide_cursor(&mut self) {
170        self.write_raw(&[0x1b, '[' as u8, '?' as u8, '2' as u8, '5' as u8, 'l' as u8]);
171    }
172
173    pub fn cursor_go_to(&mut self, row: usize, col: usize) {
174        if row == 0 && col == 0 {
175            self.write_raw(&[0x1b, '[' as u8, 'H' as u8]);
176            return;
177        }
178        self.write_raw(&[0x1b, '[' as u8]);
179        self.write_raw_str(&col.to_string());
180        self.write_raw(&[';' as u8]);
181        self.write_raw_str(&row.to_string());
182        self.write_raw(&['H' as u8]);
183    }
184
185    pub fn cursor_up(&mut self, n: isize) {
186        if n == 0 {
187            return;
188        } else if n < 0 {
189            self.cursor_down(-n);
190            return;
191        }
192        self.write_raw(&[0x1b, '[' as u8]);
193        self.write_raw_str(&n.to_string());
194        self.write_raw(&['A' as u8]);
195    }
196
197    pub fn cursor_down(&mut self, n: isize) {
198        if n == 0 {
199            return;
200        } else if n < 0 {
201            self.cursor_up(-n);
202            return;
203        }
204        self.write_raw(&[0x1b, '[' as u8]);
205        self.write_raw_str(&n.to_string());
206        self.write_raw(&['B' as u8]);
207    }
208
209    pub fn cursor_forward(&mut self, n: isize) {
210        if n == 0 {
211            return;
212        } else if n < 0 {
213            self.cursor_backward(-n);
214            return;
215        }
216        self.write_raw(&[0x1b, '[' as u8]);
217        self.write_raw_str(&n.to_string());
218        self.write_raw(&['C' as u8]);
219    }
220
221    pub fn cursor_backward(&mut self, n: isize) {
222        if n == 0 {
223            return;
224        } else if n < 0 {
225            self.cursor_forward(-n);
226            return;
227        }
228        self.write_raw(&[0x1b, '[' as u8]);
229        self.write_raw_str(&n.to_string());
230        self.write_raw(&['D' as u8]);
231    }
232
233    pub fn ask_for_cpr(&mut self) {
234        self.write_raw(&[0x1b, '[' as u8, '6' as u8, 'n' as u8]);
235    }
236
237    pub fn save_cursor(&mut self) {
238        self.write_raw(&[0x1b, '[' as u8, 's' as u8]);
239    }
240
241    pub fn pop_cursor(&mut self) {
242        self.write_raw(&[0x1b, '[' as u8, 'u' as u8]);
243    }
244}
245
246impl Writer {
247    pub fn scroll_down(&mut self) {
248        self.write_raw(&[0x1b, 'D' as u8]);
249    }
250
251    pub fn scroll_up(&mut self) {
252        self.write_raw(&[0x1b, 'M' as u8]);
253    }
254}
255
256impl Writer {
257    pub fn set_title(&mut self, title: &str) {
258        let mut title_bytes = title.as_bytes().to_vec();
259        replace_vec_inplace(&mut title_bytes, 0x13, None);
260        replace_vec_inplace(&mut title_bytes, 0x07, None);
261        self.write_raw(&[0x1b, ']' as u8, '2' as u8, ';' as u8]);
262        self.write_raw(&title_bytes);
263        self.write_raw(&[0x07]);
264    }
265
266    pub fn clear_title(&mut self) {
267        self.write_raw(&[0x1b, ']' as u8, '2' as u8, ';' as u8, 0x06]);
268    }
269}
270
271impl Writer {
272    pub fn set_color(&mut self, foreground_color: Color, background_color: Color, bold: bool) {
273        match bold {
274            true => {
275                self.set_display_attributes(
276                    foreground_color,
277                    background_color,
278                    DisplayAttribute::Bold,
279                );
280            }
281            false => {
282                self.set_display_attributes(
283                    foreground_color,
284                    background_color,
285                    DisplayAttribute::Reset,
286                );
287            }
288        }
289    }
290
291    pub fn set_display_attributes(
292        &mut self,
293        foreground_color: Color,
294        background_color: Color,
295        attribute: DisplayAttribute,
296    ) {
297        self.write_raw(&[0x1b, '[' as u8]);
298        match DISPLAY_ATTRIBUTE_PARAMETERS.get(&attribute) {
299            None => {}
300            Some(param) => {
301                self.write_raw(param);
302                self.write_raw(&[';' as u8]);
303            }
304        };
305        match FOREGROUND_ANSI_COLORS.get(&foreground_color) {
306            None => {}
307            Some(fg) => {
308                self.write_raw(fg);
309                self.write_raw(&[';' as u8]);
310            }
311        }
312        match BACKGROUND_ANSI_COLORS.get(&background_color) {
313            None => {}
314            Some(bg) => {
315                self.write_raw(bg);
316            }
317        }
318        self.write_raw(&['m' as u8]);
319    }
320
321    pub fn get_changed(&self) -> bool {
322        self.changed
323    }
324}
325
326impl Default for Writer {
327    fn default() -> Self {
328        Writer::new()
329    }
330}
331
332impl Drop for Writer {
333    fn drop(&mut self) {
334        self.erase_down();
335        self.write_raw(&['\n' as u8, '\r' as u8]);
336        self.flush().unwrap();
337    }
338}
339
340lazy_static! {
341    static ref DISPLAY_ATTRIBUTE_PARAMETERS: HashMap<DisplayAttribute, Vec<u8>> = {
342        let mut h = HashMap::new();
343        h.insert(DisplayAttribute::Reset, ['0' as u8].to_vec());
344        h.insert(DisplayAttribute::Bold, ['1' as u8].to_vec());
345        h.insert(DisplayAttribute::LowIntensity, ['2' as u8].to_vec());
346        h.insert(DisplayAttribute::Italic, ['3' as u8].to_vec());
347        h.insert(DisplayAttribute::Underline, ['4' as u8].to_vec());
348        h.insert(DisplayAttribute::Blink, ['5' as u8].to_vec());
349        h.insert(DisplayAttribute::RapidBlink, ['6' as u8].to_vec());
350        h.insert(DisplayAttribute::Reverse, ['7' as u8].to_vec());
351        h.insert(DisplayAttribute::Invisible, ['8' as u8].to_vec());
352        h.insert(DisplayAttribute::CrossedOut, ['9' as u8].to_vec());
353        h.insert(
354            DisplayAttribute::DefaultFont,
355            ['1' as u8, '0' as u8].to_vec(),
356        );
357        h
358    };
359    static ref FOREGROUND_ANSI_COLORS: HashMap<Color, Vec<u8>> = {
360        let mut h = HashMap::new();
361        h.insert(Color::Black, ['3' as u8, '0' as u8].to_vec());
362        h.insert(Color::DarkRed, ['3' as u8, '1' as u8].to_vec());
363        h.insert(Color::DarkGreen, ['3' as u8, '2' as u8].to_vec());
364        h.insert(Color::Brown, ['3' as u8, '3' as u8].to_vec());
365        h.insert(Color::DarkBlue, ['3' as u8, '4' as u8].to_vec());
366        h.insert(Color::Purple, ['3' as u8, '5' as u8].to_vec());
367        h.insert(Color::Cyan, ['3' as u8, '6' as u8].to_vec());
368        h.insert(Color::LightGray, ['3' as u8, '7' as u8].to_vec());
369        h.insert(Color::DarkGray, ['9' as u8, '0' as u8].to_vec());
370        h.insert(Color::Red, ['9' as u8, '1' as u8].to_vec());
371        h.insert(Color::Green, ['9' as u8, '2' as u8].to_vec());
372        h.insert(Color::Yellow, ['9' as u8, '3' as u8].to_vec());
373        h.insert(Color::Blue, ['9' as u8, '4' as u8].to_vec());
374        h.insert(Color::Fuchsia, ['9' as u8, '5' as u8].to_vec());
375        h.insert(Color::Turquoise, ['9' as u8, '6' as u8].to_vec());
376        h.insert(Color::White, ['9' as u8, '7' as u8].to_vec());
377        h
378    };
379    static ref BACKGROUND_ANSI_COLORS: HashMap<Color, Vec<u8>> = {
380        let mut h = HashMap::new();
381        h.insert(Color::Default, ['4' as u8, '9' as u8].to_vec());
382        h.insert(Color::Black, ['4' as u8, '0' as u8].to_vec());
383        h.insert(Color::DarkRed, ['4' as u8, '1' as u8].to_vec());
384        h.insert(Color::DarkGreen, ['4' as u8, '2' as u8].to_vec());
385        h.insert(Color::Brown, ['4' as u8, '3' as u8].to_vec());
386        h.insert(Color::DarkBlue, ['4' as u8, '4' as u8].to_vec());
387        h.insert(Color::Purple, ['4' as u8, '5' as u8].to_vec());
388        h.insert(Color::Cyan, ['4' as u8, '6' as u8].to_vec());
389        h.insert(Color::LightGray, ['4' as u8, '7' as u8].to_vec());
390        h.insert(Color::DarkGray, ['1' as u8, '0' as u8, '0' as u8].to_vec());
391        h.insert(Color::Red, ['1' as u8, '0' as u8, '1' as u8].to_vec());
392        h.insert(Color::Green, ['1' as u8, '0' as u8, '2' as u8].to_vec());
393        h.insert(Color::Yellow, ['1' as u8, '0' as u8, '3' as u8].to_vec());
394        h.insert(Color::Blue, ['1' as u8, '0' as u8, '4' as u8].to_vec());
395        h.insert(Color::Fuchsia, ['1' as u8, '0' as u8, '5' as u8].to_vec());
396        h.insert(Color::Turquoise, ['1' as u8, '0' as u8, '6' as u8].to_vec());
397        h.insert(Color::White, ['1' as u8, '0' as u8, '7' as u8].to_vec());
398        h
399    };
400}
401
402fn replace_vec<T: PartialEq + Clone>(v: &[T], to_replace: T, with: Option<T>) -> Vec<T> {
403    match with {
404        None => {
405            let mut output = vec![];
406            for val in v.iter() {
407                if *val != to_replace {
408                    output.push(val.clone());
409                }
410            }
411            output
412        }
413        Some(with) => v.iter().fold(Vec::new(), |mut v, value| {
414            let to_insert = if *value == to_replace {
415                with.clone()
416            } else {
417                value.clone()
418            };
419            v.push(to_insert);
420            v
421        }),
422    }
423}
424
425fn replace_vec_inplace<T: PartialEq + Clone>(v: &mut Vec<T>, to_replace: T, with: Option<T>) {
426    match with {
427        None => {
428            v.retain(|v| *v != to_replace);
429        }
430        Some(with) => {
431            v.iter_mut().for_each(|value| {
432                if *value == to_replace {
433                    *value = with.clone();
434                }
435            });
436        }
437    }
438}
439
440#[cfg(test)]
441mod tests {
442    use crate::writer::{replace_vec, replace_vec_inplace};
443
444    #[test]
445    fn test_replace_vec() {
446        let mut vec: Vec<u8> = vec![1, 2, 3, 4, 5];
447        let output = replace_vec(&mut vec, 1, Some(0));
448        assert_eq!(output[0], 0);
449        let output = replace_vec(&mut vec, 5, None);
450        assert_eq!(output.len(), 4);
451        assert_eq!(output[3], 4);
452    }
453
454    #[test]
455    fn test_replace_vec_inplace() {
456        let mut vec: Vec<u8> = vec![1, 2, 3, 4, 5];
457        replace_vec_inplace(&mut vec, 1, Some(0));
458        assert_eq!(vec[0], 0);
459        replace_vec_inplace(&mut vec, 5, None);
460        assert_eq!(vec.len(), 4);
461        assert_eq!(vec[3], 4);
462    }
463}