Skip to main content

csv_legacy/
writer.rs

1use std::io::Write;
2
3use crate::error::WriteError;
4
5/// Writes CSV data to a `std::io::Write` sink.
6///
7/// The writer handles auto-quoting of fields that contain the delimiter,
8/// newlines, or double-quotes. It also validates that all rows have the
9/// same number of fields (unless flexible mode is enabled).
10///
11/// Internal buffering is used to avoid many small writes to the
12/// underlying writer. Callers should not wrap the writer in a
13/// `BufWriter`.
14///
15/// # Example
16///
17/// ```no_run
18/// # use csv_legacy::Writer;
19/// let mut w = Writer::new(Vec::new());
20/// w.write_row(["name", "age", "city"]).unwrap();
21/// w.write_row(["Alice", "30", "New York, NY"]).unwrap();
22/// let result = String::from_utf8(w.into_inner().unwrap()).unwrap();
23/// assert_eq!(result, "name,age,city\r\nAlice,30,\"New York, NY\"\r\n");
24/// ```
25pub struct Writer<W: Write> {
26    writer: Option<W>,
27    delimiter: u8,
28    flexible: bool,
29    num_fields: Option<usize>,
30    row_count: usize,
31    buf: Vec<u8>,
32    field_buf: Vec<u8>,
33}
34
35impl<W: Write> Writer<W> {
36    /// Create a new writer wrapping the given output sink.
37    pub fn new(writer: W) -> Self {
38        Writer {
39            writer: Some(writer),
40            delimiter: b',',
41            flexible: false,
42            num_fields: None,
43            row_count: 0,
44            buf: Vec::with_capacity(8192),
45            field_buf: Vec::new(),
46        }
47    }
48
49    /// Set the field delimiter byte (default is `,`).
50    pub fn delimiter(&mut self, byte: u8) -> &mut Self {
51        self.delimiter = byte;
52        self
53    }
54
55    /// Enable flexible mode, which skips the field count consistency check.
56    pub fn flexible(&mut self, yes: bool) -> &mut Self {
57        self.flexible = yes;
58        self
59    }
60
61    /// Write a single row. Each element in the iterator is written as a
62    /// CSV field. Fields are auto-quoted if they contain the delimiter,
63    /// a newline (`\n` or `\r`), or a double-quote (`"`).
64    ///
65    /// Whitespace trimming is never performed. The empty byte slice `&[]`
66    /// produces an empty field.
67    ///
68    /// # Errors
69    ///
70    /// Returns [`WriteError::InconsistentFieldCount`] if the number of
71    /// fields differs from previous rows (unless flexible mode is enabled).
72    /// Returns [`WriteError::Io`] if the underlying writer fails.
73    pub fn write_row<I, T>(&mut self, row: I) -> Result<(), WriteError>
74    where
75        I: IntoIterator<Item = T>,
76        T: AsRef<[u8]>,
77    {
78        let mut field_count = 0;
79
80        for field_bytes in row {
81            if field_count > 0 {
82                self.buf.push(self.delimiter);
83            }
84            self.write_field(field_bytes.as_ref())?;
85            field_count += 1;
86        }
87
88        self.buf.extend_from_slice(b"\r\n");
89        self.row_count += 1;
90
91        // Validate field count
92        match self.num_fields {
93            Some(expected) if !self.flexible && field_count != expected => {
94                return Err(WriteError::InconsistentFieldCount {
95                    expected,
96                    found: field_count,
97                    row: self.row_count,
98                });
99            }
100            None => {
101                self.num_fields = Some(field_count);
102            }
103            _ => {}
104        }
105
106        // Flush buffer if it's large enough
107        if self.buf.len() >= 8192 {
108            self.flush()?;
109        }
110
111        Ok(())
112    }
113
114    fn write_field(&mut self, field: &[u8]) -> Result<(), WriteError> {
115        let needs_quoting = field.is_empty()
116            || field.contains(&self.delimiter)
117            || field.contains(&b'\n')
118            || field.contains(&b'\r')
119            || field.contains(&b'"');
120
121        if !needs_quoting {
122            self.buf.extend_from_slice(field);
123            return Ok(());
124        }
125
126        self.buf.push(b'"');
127        self.field_buf.clear();
128
129        for &byte in field {
130            if byte == b'"' {
131                self.field_buf.push(b'"');
132                self.field_buf.push(b'"');
133            } else {
134                self.field_buf.push(byte);
135            }
136        }
137
138        self.buf.extend_from_slice(&self.field_buf);
139        self.buf.push(b'"');
140
141        Ok(())
142    }
143
144    /// Flush the internal buffer to the underlying writer.
145    pub fn flush(&mut self) -> Result<(), WriteError> {
146        if let Some(writer) = self.writer.as_mut() {
147            writer.write_all(&self.buf)?;
148        }
149        self.buf.clear();
150        Ok(())
151    }
152
153    /// Unwrap the writer, flushing any remaining data and returning
154    /// the underlying writer.
155    pub fn into_inner(mut self) -> Result<W, WriteError> {
156        self.flush()?;
157        Ok(self.writer.take().unwrap())
158    }
159}
160
161impl<W: Write> Drop for Writer<W> {
162    fn drop(&mut self) {
163        if self.writer.is_some() {
164            let _ = self.flush();
165        }
166    }
167}