anstream/
strip.rs

1use crate::adapter::StripBytes;
2use crate::stream::AsLockedWrite;
3use crate::stream::IsTerminal;
4
5/// Only pass printable data to the inner `Write`
6#[derive(Debug)]
7pub struct StripStream<S>
8where
9    S: std::io::Write,
10{
11    raw: S,
12    state: StripBytes,
13}
14
15impl<S> StripStream<S>
16where
17    S: std::io::Write,
18{
19    /// Only pass printable data to the inner `Write`
20    #[inline]
21    pub fn new(raw: S) -> Self {
22        Self {
23            raw,
24            state: Default::default(),
25        }
26    }
27
28    /// Get the wrapped [`std::io::Write`]
29    #[inline]
30    pub fn into_inner(self) -> S {
31        self.raw
32    }
33}
34
35impl<S> StripStream<S>
36where
37    S: std::io::Write,
38    S: IsTerminal,
39{
40    #[inline]
41    pub fn is_terminal(&self) -> bool {
42        self.raw.is_terminal()
43    }
44}
45
46impl StripStream<std::io::Stdout> {
47    /// Get exclusive access to the `StripStream`
48    ///
49    /// Why?
50    /// - Faster performance when writing in a loop
51    /// - Avoid other threads interleaving output with the current thread
52    #[inline]
53    pub fn lock(self) -> StripStream<std::io::StdoutLock<'static>> {
54        StripStream {
55            raw: self.raw.lock(),
56            state: self.state,
57        }
58    }
59}
60
61impl StripStream<std::io::Stderr> {
62    /// Get exclusive access to the `StripStream`
63    ///
64    /// Why?
65    /// - Faster performance when writing in a loop
66    /// - Avoid other threads interleaving output with the current thread
67    #[inline]
68    pub fn lock(self) -> StripStream<std::io::StderrLock<'static>> {
69        StripStream {
70            raw: self.raw.lock(),
71            state: self.state,
72        }
73    }
74}
75
76impl<S> std::io::Write for StripStream<S>
77where
78    S: std::io::Write,
79    S: AsLockedWrite,
80{
81    // Must forward all calls to ensure locking happens appropriately
82    #[inline]
83    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
84        write(&mut self.raw.as_locked_write(), &mut self.state, buf)
85    }
86    #[inline]
87    fn write_vectored(&mut self, bufs: &[std::io::IoSlice<'_>]) -> std::io::Result<usize> {
88        let buf = bufs
89            .iter()
90            .find(|b| !b.is_empty())
91            .map(|b| &**b)
92            .unwrap_or(&[][..]);
93        self.write(buf)
94    }
95    // is_write_vectored: nightly only
96    #[inline]
97    fn flush(&mut self) -> std::io::Result<()> {
98        self.raw.as_locked_write().flush()
99    }
100    #[inline]
101    fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
102        write_all(&mut self.raw.as_locked_write(), &mut self.state, buf)
103    }
104    // write_all_vectored: nightly only
105    #[inline]
106    fn write_fmt(&mut self, args: std::fmt::Arguments<'_>) -> std::io::Result<()> {
107        write_fmt(&mut self.raw.as_locked_write(), &mut self.state, args)
108    }
109}
110
111fn write(
112    raw: &mut dyn std::io::Write,
113    state: &mut StripBytes,
114    buf: &[u8],
115) -> std::io::Result<usize> {
116    let initial_state = state.clone();
117
118    for printable in state.strip_next(buf) {
119        let possible = printable.len();
120        let written = raw.write(printable)?;
121        if possible != written {
122            let divergence = &printable[written..];
123            let offset = offset_to(buf, divergence);
124            let consumed = &buf[offset..];
125            *state = initial_state;
126            state.strip_next(consumed).last();
127            return Ok(offset);
128        }
129    }
130    Ok(buf.len())
131}
132
133fn write_all(
134    raw: &mut dyn std::io::Write,
135    state: &mut StripBytes,
136    buf: &[u8],
137) -> std::io::Result<()> {
138    for printable in state.strip_next(buf) {
139        raw.write_all(printable)?;
140    }
141    Ok(())
142}
143
144fn write_fmt(
145    raw: &mut dyn std::io::Write,
146    state: &mut StripBytes,
147    args: std::fmt::Arguments<'_>,
148) -> std::io::Result<()> {
149    let write_all = |buf: &[u8]| write_all(raw, state, buf);
150    crate::fmt::Adapter::new(write_all).write_fmt(args)
151}
152
153#[inline]
154fn offset_to(total: &[u8], subslice: &[u8]) -> usize {
155    let total = total.as_ptr();
156    let subslice = subslice.as_ptr();
157
158    debug_assert!(
159        total <= subslice,
160        "`Offset::offset_to` only accepts slices of `self`"
161    );
162    subslice as usize - total as usize
163}
164
165#[cfg(test)]
166mod test {
167    use super::*;
168    use proptest::prelude::*;
169    use std::io::Write as _;
170
171    proptest! {
172        #[test]
173        #[cfg_attr(miri, ignore)]  // See https://github.com/AltSysrq/proptest/issues/253
174        fn write_all_no_escapes(s in "\\PC*") {
175            let buffer = Vec::new();
176            let mut stream = StripStream::new(buffer);
177            stream.write_all(s.as_bytes()).unwrap();
178            let buffer = stream.into_inner();
179            let actual = std::str::from_utf8(buffer.as_ref()).unwrap();
180            assert_eq!(s, actual);
181        }
182
183        #[test]
184        #[cfg_attr(miri, ignore)]  // See https://github.com/AltSysrq/proptest/issues/253
185        fn write_byte_no_escapes(s in "\\PC*") {
186            let buffer = Vec::new();
187            let mut stream = StripStream::new(buffer);
188            for byte in s.as_bytes() {
189                stream.write_all(&[*byte]).unwrap();
190            }
191            let buffer = stream.into_inner();
192            let actual = std::str::from_utf8(buffer.as_ref()).unwrap();
193            assert_eq!(s, actual);
194        }
195
196        #[test]
197        #[cfg_attr(miri, ignore)]  // See https://github.com/AltSysrq/proptest/issues/253
198        fn write_all_random(s in any::<Vec<u8>>()) {
199            let buffer = Vec::new();
200            let mut stream = StripStream::new(buffer);
201            stream.write_all(s.as_slice()).unwrap();
202            let buffer = stream.into_inner();
203            if let Ok(actual) = std::str::from_utf8(buffer.as_ref()) {
204                for char in actual.chars() {
205                    assert!(!char.is_ascii() || !char.is_control() || char.is_ascii_whitespace(), "{:?} -> {:?}: {:?}", String::from_utf8_lossy(&s), actual, char);
206                }
207            }
208        }
209
210        #[test]
211        #[cfg_attr(miri, ignore)]  // See https://github.com/AltSysrq/proptest/issues/253
212        fn write_byte_random(s in any::<Vec<u8>>()) {
213            let buffer = Vec::new();
214            let mut stream = StripStream::new(buffer);
215            for byte in s.as_slice() {
216                stream.write_all(&[*byte]).unwrap();
217            }
218            let buffer = stream.into_inner();
219            if let Ok(actual) = std::str::from_utf8(buffer.as_ref()) {
220                for char in actual.chars() {
221                    assert!(!char.is_ascii() || !char.is_control() || char.is_ascii_whitespace(), "{:?} -> {:?}: {:?}", String::from_utf8_lossy(&s), actual, char);
222                }
223            }
224        }
225    }
226}
anstream/strip.rs

anstream/
strip.rs
