src/render/ansi/sgr/writer.rs

//! Writer that emits SGR escape sequences when color is enabled.
//!
//! Respects the `NO_COLOR` convention from <https://no-color.org>. Callers
//! can also force color off via the constructor argument.

use std::io::{self, Write};

use super::palette::Color;

pub struct SgrWriter<W: Write> {
    inner: W,
    enabled: bool,
}

impl<W: Write> SgrWriter<W> {
    pub fn new(inner: W, color: bool) -> Self {
        let enabled = color && std::env::var_os("NO_COLOR").is_none();
        Self { inner, enabled }
    }

    pub fn with_env(inner: W) -> Self {
        let enabled = std::env::var_os("NO_COLOR").is_none();
        Self { inner, enabled }
    }

    pub fn raw(&mut self) -> &mut W {
        &mut self.inner
    }

    pub fn enabled(&self) -> bool {
        self.enabled
    }

    pub fn with_fg<F>(&mut self, color: Color, f: F) -> io::Result<()>
    where
        F: FnOnce(&mut W) -> io::Result<()>,
    {
        if self.enabled {
            write!(self.inner, "\x1b[38;5;{}m", color)?;
        }
        f(&mut self.inner)?;
        if self.enabled {
            self.inner.write_all(b"\x1b[0m")?;
        }
        Ok(())
    }

    pub fn bold<F>(&mut self, f: F) -> io::Result<()>
    where
        F: FnOnce(&mut W) -> io::Result<()>,
    {
        if self.enabled {
            self.inner.write_all(b"\x1b[1m")?;
        }
        f(&mut self.inner)?;
        if self.enabled {
            self.inner.write_all(b"\x1b[22m")?;
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn disabled_writer_omits_escape() {
        let mut buf = Vec::new();
        {
            let mut w = SgrWriter::new(&mut buf, false);
            w.with_fg(10, |inner| inner.write_all(b"hi")).unwrap();
        }
        assert_eq!(buf, b"hi");
    }

    #[test]
    fn enabled_writer_brackets_with_escape() {
        let mut buf = Vec::new();
        {
            let mut w = SgrWriter::new(&mut buf, true);
            std::env::remove_var("NO_COLOR");
            w.with_fg(10, |inner| inner.write_all(b"hi")).unwrap();
        }
        assert!(buf.starts_with(b"\x1b["));
        assert!(buf.ends_with(b"\x1b[0m"));
    }
}