tracing-subscriber-multi 0.1.0

use std::io::Write;

/// Strip ANSI control sequences.
///
/// Note, this only supports stripping CSI sequences and is not intended for
/// purposes other than with a logger.
pub struct AnsiStripper<W: Write> {
    inner: W,
    state: AnsiStateMachine,
}

enum AnsiStateMachine {
    ScanningForESC,
    ExpectingType,
    WaitingForEnd,
}

impl<W: Write> AnsiStripper<W> {
    /// Create a new writer that strips ANSI control sequences.
    pub fn new(writer: W) -> Self {
        Self {
            inner: writer,
            state: AnsiStateMachine::ScanningForESC,
        }
    }
}

impl<W: Write> Write for AnsiStripper<W> {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let mut cleaned_buffer = vec![];
        for &c in buf {
            match &mut self.state {
                AnsiStateMachine::ScanningForESC => {
                    if c == 0x1b {
                        // ANSI sequence start
                        self.state = AnsiStateMachine::ExpectingType;
                    } else {
                        // Regular data
                        cleaned_buffer.push(c);
                    }
                }
                AnsiStateMachine::ExpectingType => {
                    if (0x40..=0x5f).contains(&c) {
                        // Fe type
                        self.state = AnsiStateMachine::WaitingForEnd;
                    } else if (0x60..=0x7e).contains(&c) {
                        // Fs type
                        tracing::warn!("The writer was asked to strip unsupported ANSI sequences!");
                    } else if (0x30..=0x3f).contains(&c) {
                        // Fp type
                        tracing::warn!("The writer was asked to strip unsupported ANSI sequences!");
                    } else if (0x20..=0x2f).contains(&c) {
                        // nF type
                        tracing::warn!("The writer was asked to strip unsupported ANSI sequences!");
                    }
                }
                AnsiStateMachine::WaitingForEnd => {
                    if (0x40..=0x7e).contains(&c) {
                        // Fe sequence, skip 1 char
                        self.state = AnsiStateMachine::ScanningForESC;
                    }
                }
            }
        }

        self.inner.write(&cleaned_buffer)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.inner.flush()
    }
}