indent_write/
io.rs

1use std::io;
2
3use super::Inspect;
4
5#[derive(Debug, Copy, Clone)]
6enum IndentState<'a> {
7    // We are currently writing a line. Forward writes until the end of the
8    // line.
9    MidLine,
10
11    // An indent has been requested. Write empty lines, then write an indent
12    // before the next non empty line.
13    NeedIndent,
14
15    // We are currently writing an indent.
16    WritingIndent(&'a [u8]),
17}
18
19use IndentState::*;
20
21/// Adapter for writers to indent each line
22///
23/// An `IndentWriter` adapts an [`io::Write`] object to insert an indent before
24/// each non-empty line. Specifically, this means it will insert an indent
25/// between each newline when followed by a non-newline.
26///
27/// These writers can be nested to provide increasing levels of indentation.
28///
29/// # Example
30///
31/// ```
32/// # use std::io::Write;
33/// use indent_write::io::IndentWriter;
34///
35/// let output = Vec::new();
36///
37/// let mut indented = IndentWriter::new("\t", output);
38///
39/// // Lines will be indented
40/// write!(indented, "Line 1\nLine 2\n");
41///
42/// // Empty lines will not be indented
43/// write!(indented, "\n\nLine 3\n\n");
44///
45/// assert_eq!(indented.get_ref(), b"\tLine 1\n\tLine 2\n\n\n\tLine 3\n\n");
46/// ```
47#[derive(Debug, Clone)]
48pub struct IndentWriter<'i, W> {
49    writer: W,
50    indent: &'i str,
51    state: IndentState<'i>,
52}
53
54impl<'i, W: io::Write> IndentWriter<'i, W> {
55    /// Create a new [`IndentWriter`].
56    pub fn new(indent: &'i str, writer: W) -> Self {
57        Self {
58            writer,
59            indent,
60            state: NeedIndent,
61        }
62    }
63
64    /// Create a new [`IndentWriter`] which will not add an indent to the first
65    /// written line.
66    ///
67    /// # Example
68    ///
69    /// ```
70    /// # use std::io::Write;
71    /// use indent_write::io::IndentWriter;
72    ///
73    /// let mut buffer = Vec::new();
74    /// let mut writer = IndentWriter::new_skip_initial("    ", &mut buffer);
75    ///
76    /// writeln!(writer, "Line 1").unwrap();
77    /// writeln!(writer, "Line 2").unwrap();
78    /// writeln!(writer, "Line 3").unwrap();
79    ///
80    /// assert_eq!(buffer, b"Line 1\n    Line 2\n    Line 3\n")
81    /// ```
82    #[inline]
83    pub fn new_skip_initial(indent: &'i str, writer: W) -> Self {
84        Self {
85            writer,
86            indent,
87            state: MidLine,
88        }
89    }
90
91    /// Extract the writer from the [`IndentWriter`], discarding any in-progress
92    /// indent state.
93    #[inline]
94    pub fn into_inner(self) -> W {
95        self.writer
96    }
97
98    /// Get a reference to the wrapped writer
99    #[inline]
100    pub fn get_ref(&self) -> &W {
101        &self.writer
102    }
103
104    /// Get the string being used as an indent for each line
105    #[inline]
106    pub fn indent(&self) -> &'i str {
107        self.indent
108    }
109}
110
111impl<'i, W: io::Write> io::Write for IndentWriter<'i, W> {
112    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
113        loop {
114            match self.state {
115                // We're currently writing a line. Scan for the end of the line.
116                IndentState::MidLine => match buf.iter().position(|&b| b == b'\n') {
117                    // No newlines in the input buffer, so write the entire thing.
118                    None => break self.writer.write(buf),
119
120                    // We are at a newline presently. Request an indent be
121                    // written at the front of the next non-empty line, then
122                    // continue looping (since we haven't yet attempted to
123                    // write user data).
124                    Some(0) => self.state = NeedIndent,
125
126                    // There's an upcoming newline. Write out the remainder of
127                    // this line, plus its newline. If the entire line was
128                    // written, request an indent on the subsequent call to
129                    // write.
130                    Some(len) => {
131                        break self.writer.write(&buf[..len + 1]).inspect(|&n| {
132                            if n >= len {
133                                self.state = NeedIndent;
134                            }
135                        })
136                    }
137                },
138
139                // We need an indent. Scan for the next non-empty line.
140                IndentState::NeedIndent => match buf.iter().position(|&b| b != b'\n') {
141                    // No non-empty lines in the input buffer, so write the entire thing
142                    None => break self.writer.write(buf),
143
144                    // We are at the beginning of a non-empty line presently.
145                    // Begin inserting an indent now, then continue looping
146                    // (since we haven't yet attempted to write user data)
147                    Some(0) => self.state = WritingIndent(self.indent.as_bytes()),
148
149                    // There's an upcoming non-empty line. Write out the
150                    // remainder of the empty lines. If all the empty lines
151                    // were written, force an indent on the subsequent call to
152                    // write.
153                    Some(len) => {
154                        break self.writer.write(&buf[..len]).inspect(|&n| {
155                            if n >= len {
156                                self.state = IndentState::WritingIndent(self.indent.as_bytes())
157                            }
158                        })
159                    }
160                },
161
162                // We are writing an indent unconditionally. If we're in this
163                // state, the input buffer is known to be the start of a non-
164                // empty line.
165                IndentState::WritingIndent(indent) => match self.writer.write(indent)? {
166                    // We successfully wrote the entire indent. Continue with
167                    // writing the input buffer.
168                    n if n >= indent.len() => self.state = MidLine,
169
170                    // Eof; stop work immediately
171                    0 => break Ok(0),
172
173                    // Only a part of the indent was written. Continue
174                    // trying to write the rest of it, but update our state
175                    // to keep it consistent in case the next write is an
176                    // error
177                    n => self.state = WritingIndent(&indent[n..]),
178                },
179            }
180        }
181    }
182
183    fn flush(&mut self) -> io::Result<()> {
184        // If we're currently in the middle of writing an indent, flush it
185        while let WritingIndent(ref mut indent) = self.state {
186            match self.writer.write(*indent)? {
187                // We wrote the entire indent. Proceed with the flush
188                len if len >= indent.len() => self.state = MidLine,
189
190                // EoF; return an error
191                0 => return Err(io::ErrorKind::WriteZero.into()),
192
193                // Partial write, continue writing.
194                len => *indent = &indent[len..],
195            }
196        }
197
198        self.writer.flush()
199    }
200}