Skip to main content

csv_legacy2/
writer.rs

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