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

anstream/
strip.rs
