serde_generate/
indent.rs

1// Copyright (c) Facebook, Inc. and its affiliates
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use std::io::{Result, Write};
5
6#[derive(Clone, Copy)]
7pub enum IndentConfig {
8    Tab,
9    Space(usize),
10}
11
12pub struct IndentedWriter<T> {
13    out: T,
14    indentation: Vec<u8>,
15    config: IndentConfig,
16    at_begining_of_line: bool,
17}
18
19impl<T> IndentedWriter<T> {
20    pub fn new(out: T, config: IndentConfig) -> Self {
21        Self {
22            out,
23            indentation: Vec::new(),
24            config,
25            at_begining_of_line: true,
26        }
27    }
28
29    pub fn indent(&mut self) {
30        match self.config {
31            IndentConfig::Tab => {
32                self.indentation.push(b'\t');
33            }
34            IndentConfig::Space(n) => {
35                self.indentation.resize(self.indentation.len() + n, b' ');
36            }
37        }
38    }
39
40    pub fn unindent(&mut self) {
41        match self.config {
42            IndentConfig::Tab => {
43                self.indentation.pop();
44            }
45            IndentConfig::Space(n) => {
46                self.indentation
47                    .truncate(self.indentation.len().saturating_sub(n));
48            }
49        }
50    }
51}
52
53impl<T: Write> Write for IndentedWriter<T> {
54    fn write(&mut self, mut buf: &[u8]) -> Result<usize> {
55        let mut bytes_written = 0;
56
57        while !buf.is_empty() {
58            let (before_newline, has_newline, after_newline) =
59                if let Some(idx) = buf.iter().position(|&b| b == b'\n') {
60                    (&buf[..idx], true, &buf[idx + 1..])
61                } else {
62                    (buf, false, &buf[buf.len()..])
63                };
64
65            if self.at_begining_of_line && !before_newline.is_empty() {
66                self.out.write_all(&self.indentation)?;
67                self.at_begining_of_line = false;
68            }
69
70            self.out.write_all(before_newline)?;
71            bytes_written += before_newline.len();
72
73            if has_newline {
74                self.out.write_all(b"\n")?;
75                bytes_written += 1;
76                self.at_begining_of_line = true;
77            }
78
79            buf = after_newline;
80        }
81
82        Ok(bytes_written)
83    }
84
85    fn flush(&mut self) -> Result<()> {
86        self.out.flush()
87    }
88}
89
90#[cfg(test)]
91mod test {
92    use super::*;
93
94    #[test]
95    fn basic() -> Result<()> {
96        let mut buffer: Vec<u8> = Vec::new();
97
98        let mut out = IndentedWriter::new(&mut buffer, IndentConfig::Space(2));
99
100        writeln!(out, "foo")?;
101        out.indent();
102        writeln!(out, "bar")?;
103        writeln!(out)?;
104        writeln!(out, "bar")?;
105        out.indent();
106        writeln!(out, "foobar")?;
107        writeln!(out)?;
108        writeln!(out, "foobar")?;
109        out.unindent();
110        out.unindent();
111        writeln!(out, "foo")?;
112
113        let expect: &[u8] = b"\
114foo
115  bar
116
117  bar
118    foobar
119
120    foobar
121foo
122";
123        assert_eq!(buffer, expect);
124
125        Ok(())
126    }
127}
serde_generate/indent.rs

serde_generate/
indent.rs
