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