diff --git a/crates/anstyle-stream/src/auto.rs b/crates/anstyle-stream/src/auto.rs new file mode 100644 index 00000000..9429cdd4 --- /dev/null +++ b/crates/anstyle-stream/src/auto.rs @@ -0,0 +1,129 @@ +use crate::Lockable; +use crate::RawStream; +use crate::StripStream; + +/// [`std::io::Write`] that adapts ANSI escape codes to the underlying `Write`s capabilities +pub struct AutoStream { + inner: StreamInner, +} + +enum StreamInner { + PassThrough(S), + Strip(StripStream), +} + +impl AutoStream +where + S: RawStream, +{ + /// Force ANSI escape codes to be passed through as-is, no matter what the inner `Write` + /// supports. + #[inline] + pub fn always_ansi(raw: S) -> Self { + #[cfg(feature = "auto")] + { + if raw.is_terminal() { + let _ = concolor_query::windows::enable_ansi_colors(); + } + } + Self::always_ansi_(raw) + } + + fn always_ansi_(raw: S) -> Self { + let inner = StreamInner::PassThrough(raw); + AutoStream { inner } + } + + /// Only pass printable data to the inner `Write`. + #[inline] + pub fn never(raw: S) -> Self { + let inner = StreamInner::Strip(StripStream::new(raw)); + AutoStream { inner } + } +} + +impl AutoStream +where + S: Lockable, +{ + /// Get exclusive access to the `AutoStream` + /// + /// Why? + /// - Faster performance when writing in a loop + /// - Avoid other threads interleaving output with the current thread + #[inline] + pub fn lock(self) -> ::Locked { + let inner = match self.inner { + StreamInner::PassThrough(w) => StreamInner::PassThrough(w.lock()), + StreamInner::Strip(w) => StreamInner::Strip(w.lock()), + }; + AutoStream { inner } + } +} + +#[cfg(feature = "auto")] +impl AutoStream +where + S: RawStream, +{ + #[cfg(feature = "auto")] + #[inline] + pub(crate) fn auto(raw: S) -> Self { + if raw.is_terminal() { + if concolor_query::windows::enable_ansi_colors().unwrap_or(true) { + Self::always_ansi_(raw) + } else { + Self::never(raw) + } + } else { + Self::never(raw) + } + } +} + +impl std::io::Write for AutoStream +where + S: RawStream, +{ + #[inline] + fn write(&mut self, buf: &[u8]) -> std::io::Result { + match &mut self.inner { + StreamInner::PassThrough(w) => w.write(buf), + StreamInner::Strip(w) => w.write(buf), + } + } + + #[inline] + fn flush(&mut self) -> std::io::Result<()> { + match &mut self.inner { + StreamInner::PassThrough(w) => w.flush(), + StreamInner::Strip(w) => w.flush(), + } + } + + // Provide explicit implementations of trait methods + // - To reduce bookkeeping + // - Avoid acquiring / releasing locks in a loop + + #[inline] + fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> { + match &mut self.inner { + StreamInner::PassThrough(w) => w.write_all(buf), + StreamInner::Strip(w) => w.write_all(buf), + } + } + + // Not bothering with `write_fmt` as it just calls `write_all` +} + +impl Lockable for AutoStream +where + W: Lockable, +{ + type Locked = AutoStream<::Locked>; + + #[inline] + fn lock(self) -> Self::Locked { + self.lock() + } +} diff --git a/crates/anstyle-stream/src/buffer.rs b/crates/anstyle-stream/src/buffer.rs new file mode 100644 index 00000000..4afa715f --- /dev/null +++ b/crates/anstyle-stream/src/buffer.rs @@ -0,0 +1,48 @@ +/// In-memory [`RawStream`][crate::RawStream] +#[derive(Clone, Default, Debug, PartialEq, Eq)] +pub struct Buffer(Vec); + +impl Buffer { + #[inline] + pub fn new() -> Self { + Default::default() + } + + #[inline] + pub fn with_capacity(capacity: usize) -> Self { + Self(Vec::with_capacity(capacity)) + } + + #[inline] + pub fn as_bytes(&self) -> &[u8] { + &self.0 + } +} + +impl AsRef<[u8]> for Buffer { + #[inline] + fn as_ref(&self) -> &[u8] { + self.as_bytes() + } +} + +#[cfg(feature = "auto")] +impl is_terminal::IsTerminal for Buffer { + #[inline] + fn is_terminal(&self) -> bool { + false + } +} + +impl std::io::Write for Buffer { + #[inline] + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.0.extend(buf); + Ok(buf.len()) + } + + #[inline] + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} diff --git a/crates/anstyle-stream/src/lib.rs b/crates/anstyle-stream/src/lib.rs index 0dff2e38..1661ce26 100644 --- a/crates/anstyle-stream/src/lib.rs +++ b/crates/anstyle-stream/src/lib.rs @@ -1,6 +1,6 @@ //! **Auto-adapting [`stdout`] / [`stderr`] streams** //! -//! [`Stream`] always accepts [ANSI escape codes](https://en.wikipedia.org/wiki/ANSI_escape_code), +//! [`AutoStream`] always accepts [ANSI escape codes](https://en.wikipedia.org/wiki/ANSI_escape_code), //! adapting to the user's terminal's capabilities. //! //! Benefits @@ -27,269 +27,37 @@ #![cfg_attr(docsrs, feature(doc_auto_cfg))] pub mod adapter; +mod buffer; #[macro_use] mod macros; +mod auto; +mod lockable; +mod raw; +mod strip; + +pub use auto::AutoStream; +pub use lockable::Lockable; +pub use raw::RawStream; +pub use strip::StripStream; + +pub use buffer::Buffer; /// Create an ANSI escape code compatible stdout /// -/// **Note:** Call [`Stream::lock`] in loops to avoid the performance hit of acquiring/releasing +/// **Note:** Call [`AutoStream::lock`] in loops to avoid the performance hit of acquiring/releasing /// from the implicit locking in each [`std::io::Write`] call #[cfg(feature = "auto")] -pub fn stdout() -> Stream { +pub fn stdout() -> AutoStream { let stdout = std::io::stdout(); - Stream::auto(stdout) + AutoStream::auto(stdout) } /// Create an ANSI escape code compatible stderr /// -/// **Note:** Call [`Stream::lock`] in loops to avoid the performance hit of acquiring/releasing +/// **Note:** Call [`AutoStream::lock`] in loops to avoid the performance hit of acquiring/releasing /// from the implicit locking in each [`std::io::Write`] call #[cfg(feature = "auto")] -pub fn stderr() -> Stream { +pub fn stderr() -> AutoStream { let stderr = std::io::stderr(); - Stream::auto(stderr) -} - -/// Explicitly lock a [`std::io::Write`]able -pub trait Lockable { - type Locked; - - /// Get exclusive access to the `Stream` - /// - /// Why? - /// - Faster performance when writing in a loop - /// - Avoid other threads interleaving output with the current thread - fn lock(self) -> Self::Locked; -} - -impl Lockable for std::io::Stdout { - type Locked = std::io::StdoutLock<'static>; - - #[inline] - fn lock(self) -> Self::Locked { - #[allow(clippy::needless_borrow)] // Its needed to avoid recursion - (&self).lock() - } -} - -impl Lockable for std::io::Stderr { - type Locked = std::io::StderrLock<'static>; - - #[inline] - fn lock(self) -> Self::Locked { - #[allow(clippy::needless_borrow)] // Its needed to avoid recursion - (&self).lock() - } -} - -/// [`std::io::Write`] that adapts ANSI escape codes to the underlying `Write`s capabilities -pub struct Stream { - write: StreamInner, -} - -enum StreamInner { - PassThrough(W), - Strip(StripStream), -} - -impl Stream -where - W: std::io::Write, -{ - /// Force ANSI escape codes to be passed through as-is, no matter what the inner `Write` - /// supports. - #[inline] - pub fn always_ansi(write: W) -> Self { - #[cfg(feature = "auto")] - let _ = concolor_query::windows::enable_ansi_colors(); - Self::always_ansi_(write) - } - - fn always_ansi_(write: W) -> Self { - let write = StreamInner::PassThrough(write); - Stream { write } - } - - /// Only pass printable data to the inner `Write`. - #[inline] - pub fn never(write: W) -> Self { - let write = StreamInner::Strip(StripStream::new(write)); - Stream { write } - } -} - -impl Stream -where - W: Lockable, -{ - /// Get exclusive access to the `Stream` - /// - /// Why? - /// - Faster performance when writing in a loop - /// - Avoid other threads interleaving output with the current thread - #[inline] - pub fn lock(self) -> ::Locked { - let write = match self.write { - StreamInner::PassThrough(w) => StreamInner::PassThrough(w.lock()), - StreamInner::Strip(w) => StreamInner::Strip(w.lock()), - }; - Stream { write } - } -} - -#[cfg(feature = "auto")] -impl Stream -where - W: std::io::Write + is_terminal::IsTerminal, -{ - #[cfg(feature = "auto")] - #[inline] - fn auto(write: W) -> Self { - if write.is_terminal() { - if concolor_query::windows::enable_ansi_colors().unwrap_or(true) { - Self::always_ansi_(write) - } else { - Self::never(write) - } - } else { - Self::never(write) - } - } -} - -impl std::io::Write for Stream -where - W: std::io::Write, -{ - #[inline] - fn write(&mut self, buf: &[u8]) -> std::io::Result { - match &mut self.write { - StreamInner::PassThrough(w) => w.write(buf), - StreamInner::Strip(w) => w.write(buf), - } - } - - #[inline] - fn flush(&mut self) -> std::io::Result<()> { - match &mut self.write { - StreamInner::PassThrough(w) => w.flush(), - StreamInner::Strip(w) => w.flush(), - } - } - - // Provide explicit implementations of trait methods - // - To reduce bookkeeping - // - Avoid acquiring / releasing locks in a loop - - #[inline] - fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> { - match &mut self.write { - StreamInner::PassThrough(w) => w.write_all(buf), - StreamInner::Strip(w) => w.write_all(buf), - } - } - - // Not bothering with `write_fmt` as it just calls `write_all` -} - -impl Lockable for Stream -where - W: Lockable, -{ - type Locked = Stream<::Locked>; - - #[inline] - fn lock(self) -> Self::Locked { - self.lock() - } -} - -/// Only pass printable data to the inner `Write` -pub struct StripStream { - write: W, - state: adapter::StripBytes, -} - -impl StripStream -where - W: std::io::Write, -{ - /// Only pass printable data to the inner `Write` - #[inline] - pub fn new(write: W) -> Self { - Self { - write, - state: Default::default(), - } - } -} - -impl std::io::Write for StripStream -where - W: std::io::Write, -{ - fn write(&mut self, buf: &[u8]) -> std::io::Result { - let initial_state = self.state.clone(); - - let mut written = 0; - let mut possible = 0; - for printable in self.state.strip_next(buf) { - possible += printable.len(); - written += self.write.write(printable)?; - if possible != written { - let divergence = &printable[written..]; - let offset = offset_to(buf, divergence); - let consumed = &buf[offset..]; - self.state = initial_state; - self.state.strip_next(consumed).last(); - break; - } - } - Ok(written) - } - #[inline] - fn flush(&mut self) -> std::io::Result<()> { - self.write.flush() - } - - // Provide explicit implementations of trait methods - // - To reduce bookkeeping - // - Avoid acquiring / releasing locks in a loop - - #[inline] - fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> { - for printable in self.state.strip_next(buf) { - self.write.write_all(printable)?; - } - Ok(()) - } - - // Not bothering with `write_fmt` as it just calls `write_all` -} - -#[inline] -fn offset_to(total: &[u8], subslice: &[u8]) -> usize { - let total = total.as_ptr(); - let subslice = subslice.as_ptr(); - - debug_assert!( - total <= subslice, - "`Offset::offset_to` only accepts slices of `self`" - ); - subslice as usize - total as usize -} - -impl Lockable for StripStream -where - W: Lockable, -{ - type Locked = StripStream<::Locked>; - - #[inline] - fn lock(self) -> Self::Locked { - Self::Locked { - write: self.write.lock(), - state: self.state, - } - } + AutoStream::auto(stderr) } diff --git a/crates/anstyle-stream/src/lockable.rs b/crates/anstyle-stream/src/lockable.rs new file mode 100644 index 00000000..51568f84 --- /dev/null +++ b/crates/anstyle-stream/src/lockable.rs @@ -0,0 +1,31 @@ +/// Explicitly lock a [`std::io::Write`]able +pub trait Lockable { + type Locked; + + /// Get exclusive access to the `AutoStream` + /// + /// Why? + /// - Faster performance when writing in a loop + /// - Avoid other threads interleaving output with the current thread + fn lock(self) -> Self::Locked; +} + +impl Lockable for std::io::Stdout { + type Locked = std::io::StdoutLock<'static>; + + #[inline] + fn lock(self) -> Self::Locked { + #[allow(clippy::needless_borrow)] // Its needed to avoid recursion + (&self).lock() + } +} + +impl Lockable for std::io::Stderr { + type Locked = std::io::StderrLock<'static>; + + #[inline] + fn lock(self) -> Self::Locked { + #[allow(clippy::needless_borrow)] // Its needed to avoid recursion + (&self).lock() + } +} diff --git a/crates/anstyle-stream/src/macros.rs b/crates/anstyle-stream/src/macros.rs index be7df8b6..4597f8e5 100644 --- a/crates/anstyle-stream/src/macros.rs +++ b/crates/anstyle-stream/src/macros.rs @@ -9,7 +9,7 @@ /// /// **NOTE:** The `print!` macro will lock the standard output on each call. If you call /// `print!` within a hot loop, this behavior may be the bottleneck of the loop. -/// To avoid this, lock stdout with [`Stream::lock`][crate::Stream::lock]: +/// To avoid this, lock stdout with [`AutoStream::lock`][crate::AutoStream::lock]: /// ``` /// # #[cfg(feature = "auto")] { /// use std::io::Write as _; @@ -78,7 +78,7 @@ macro_rules! print { /// /// **NOTE:** The `println!` macro will lock the standard output on each call. If you call /// `println!` within a hot loop, this behavior may be the bottleneck of the loop. -/// To avoid this, lock stdout with [`Stream::lock`][crate::Stream::lock]: +/// To avoid this, lock stdout with [`AutoStream::lock`][crate::AutoStream::lock]: /// ``` /// # #[cfg(feature = "auto")] { /// use std::io::Write as _; diff --git a/crates/anstyle-stream/src/raw.rs b/crates/anstyle-stream/src/raw.rs new file mode 100644 index 00000000..a7ea64bd --- /dev/null +++ b/crates/anstyle-stream/src/raw.rs @@ -0,0 +1,29 @@ +#[cfg(not(feature = "auto"))] +pub trait RawStream: std::io::Write + private::Sealed {} + +#[cfg(feature = "auto")] +pub trait RawStream: std::io::Write + is_terminal::IsTerminal + private::Sealed {} + +impl RawStream for std::io::Stdout {} + +impl RawStream for std::io::StdoutLock<'static> {} + +impl RawStream for std::io::Stderr {} + +impl RawStream for std::io::StderrLock<'static> {} + +impl RawStream for crate::Buffer {} + +mod private { + pub trait Sealed {} + + impl Sealed for std::io::Stdout {} + + impl Sealed for std::io::StdoutLock<'static> {} + + impl Sealed for std::io::Stderr {} + + impl Sealed for std::io::StderrLock<'static> {} + + impl Sealed for crate::Buffer {} +} diff --git a/crates/anstyle-stream/src/strip.rs b/crates/anstyle-stream/src/strip.rs new file mode 100644 index 00000000..49bb9b05 --- /dev/null +++ b/crates/anstyle-stream/src/strip.rs @@ -0,0 +1,93 @@ +use crate::adapter::StripBytes; +use crate::Lockable; +use crate::RawStream; + +/// Only pass printable data to the inner `Write` +pub struct StripStream { + raw: S, + state: StripBytes, +} + +impl StripStream +where + S: RawStream, +{ + /// Only pass printable data to the inner `Write` + #[inline] + pub fn new(raw: S) -> Self { + Self { + raw, + state: Default::default(), + } + } +} + +impl std::io::Write for StripStream +where + S: RawStream, +{ + fn write(&mut self, buf: &[u8]) -> std::io::Result { + let initial_state = self.state.clone(); + + let mut written = 0; + let mut possible = 0; + for printable in self.state.strip_next(buf) { + possible += printable.len(); + written += self.raw.write(printable)?; + if possible != written { + let divergence = &printable[written..]; + let offset = offset_to(buf, divergence); + let consumed = &buf[offset..]; + self.state = initial_state; + self.state.strip_next(consumed).last(); + break; + } + } + Ok(written) + } + #[inline] + fn flush(&mut self) -> std::io::Result<()> { + self.raw.flush() + } + + // Provide explicit implementations of trait methods + // - To reduce bookkeeping + // - Avoid acquiring / releasing locks in a loop + + #[inline] + fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> { + for printable in self.state.strip_next(buf) { + self.raw.write_all(printable)?; + } + Ok(()) + } + + // Not bothering with `write_fmt` as it just calls `write_all` +} + +#[inline] +fn offset_to(total: &[u8], subslice: &[u8]) -> usize { + let total = total.as_ptr(); + let subslice = subslice.as_ptr(); + + debug_assert!( + total <= subslice, + "`Offset::offset_to` only accepts slices of `self`" + ); + subslice as usize - total as usize +} + +impl Lockable for StripStream +where + S: Lockable, +{ + type Locked = StripStream<::Locked>; + + #[inline] + fn lock(self) -> Self::Locked { + Self::Locked { + raw: self.raw.lock(), + state: self.state, + } + } +}