1 | //! Insert line breaks between written buffers when they would overflow the line length. |
2 | use std::io; |
3 | |
4 | // The pnm standard says to insert line breaks after 70 characters. Assumes that no line breaks |
5 | // are actually written. We have to be careful to fully commit buffers or not commit them at all, |
6 | // otherwise we might insert a newline in the middle of a token. |
7 | pub(crate) struct AutoBreak<W: io::Write> { |
8 | wrapped: W, |
9 | line_capacity: usize, |
10 | line: Vec<u8>, |
11 | has_newline: bool, |
12 | panicked: bool, // see https://github.com/rust-lang/rust/issues/30888 |
13 | } |
14 | |
15 | impl<W: io::Write> AutoBreak<W> { |
16 | pub(crate) fn new(writer: W, line_capacity: usize) -> Self { |
17 | AutoBreak { |
18 | wrapped: writer, |
19 | line_capacity, |
20 | line: Vec::with_capacity(line_capacity + 1), |
21 | has_newline: false, |
22 | panicked: false, |
23 | } |
24 | } |
25 | |
26 | fn flush_buf(&mut self) -> io::Result<()> { |
27 | // from BufWriter |
28 | let mut written = 0; |
29 | let len = self.line.len(); |
30 | let mut ret = Ok(()); |
31 | while written < len { |
32 | self.panicked = true; |
33 | let r = self.wrapped.write(&self.line[written..]); |
34 | self.panicked = false; |
35 | match r { |
36 | Ok(0) => { |
37 | ret = Err(io::Error::new( |
38 | io::ErrorKind::WriteZero, |
39 | "failed to write the buffered data" , |
40 | )); |
41 | break; |
42 | } |
43 | Ok(n) => written += n, |
44 | Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {} |
45 | Err(e) => { |
46 | ret = Err(e); |
47 | break; |
48 | } |
49 | } |
50 | } |
51 | if written > 0 { |
52 | self.line.drain(..written); |
53 | } |
54 | ret |
55 | } |
56 | } |
57 | |
58 | impl<W: io::Write> io::Write for AutoBreak<W> { |
59 | fn write(&mut self, buffer: &[u8]) -> io::Result<usize> { |
60 | if self.has_newline { |
61 | self.flush()?; |
62 | self.has_newline = false; |
63 | } |
64 | |
65 | if !self.line.is_empty() && self.line.len() + buffer.len() > self.line_capacity { |
66 | self.line.push(b' \n' ); |
67 | self.has_newline = true; |
68 | self.flush()?; |
69 | self.has_newline = false; |
70 | } |
71 | |
72 | self.line.extend_from_slice(buffer); |
73 | Ok(buffer.len()) |
74 | } |
75 | |
76 | fn flush(&mut self) -> io::Result<()> { |
77 | self.flush_buf()?; |
78 | self.wrapped.flush() |
79 | } |
80 | } |
81 | |
82 | impl<W: io::Write> Drop for AutoBreak<W> { |
83 | fn drop(&mut self) { |
84 | if !self.panicked { |
85 | let _r: Result<(), Error> = self.flush_buf(); |
86 | // internal writer flushed automatically by Drop |
87 | } |
88 | } |
89 | } |
90 | |
91 | #[cfg (test)] |
92 | mod tests { |
93 | use super::*; |
94 | use std::io::Write; |
95 | |
96 | #[test ] |
97 | fn test_aligned_writes() { |
98 | let mut output = Vec::new(); |
99 | |
100 | { |
101 | let mut writer = AutoBreak::new(&mut output, 10); |
102 | writer.write_all(b"0123456789" ).unwrap(); |
103 | writer.write_all(b"0123456789" ).unwrap(); |
104 | } |
105 | |
106 | assert_eq!(output.as_slice(), b"0123456789 \n0123456789" ); |
107 | } |
108 | |
109 | #[test ] |
110 | fn test_greater_writes() { |
111 | let mut output = Vec::new(); |
112 | |
113 | { |
114 | let mut writer = AutoBreak::new(&mut output, 10); |
115 | writer.write_all(b"012" ).unwrap(); |
116 | writer.write_all(b"345" ).unwrap(); |
117 | writer.write_all(b"0123456789" ).unwrap(); |
118 | writer.write_all(b"012345678910" ).unwrap(); |
119 | writer.write_all(b"_" ).unwrap(); |
120 | } |
121 | |
122 | assert_eq!(output.as_slice(), b"012345 \n0123456789 \n012345678910 \n_" ); |
123 | } |
124 | } |
125 | |