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}