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: StripBytes = state.clone();
117
118 for printable: &[u8] in state.strip_next(bytes:buf) {
119 let possible: usize = printable.len();
120 let written: usize = raw.write(buf:printable)?;
121 if possible != written {
122 let divergence: &[u8] = &printable[written..];
123 let offset: usize = offset_to(total:buf, subslice:divergence);
124 let consumed: &[u8] = &buf[offset..];
125 *state = initial_state;
126 state.strip_next(bytes: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: &[u8] in state.strip_next(bytes:buf) {
139 raw.write_all(buf: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: impl FnMut(&[u8]) -> Result<…, …> = |buf: &[u8]| write_all(raw, state, buf);
150 crate::fmt::Adapter::new(writer:write_all).write_fmt(args)
151}
152
153#[inline]
154fn offset_to(total: &[u8], subslice: &[u8]) -> usize {
155 let total: *const u8 = total.as_ptr();
156 let subslice: *const u8 = 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}
227