csv/writer.rs
1use alloc::{
2 string::{String, ToString},
3 vec::Vec,
4};
5
6use crate::error::WriteError;
7
8/// Custom write abstraction that allows `Writer` to work without `std::io::Write`.
9///
10/// This trait is automatically implemented for `Vec<u8>` and, when the `std`
11/// feature is enabled, for any `std::io::Write` type.
12pub trait Write {
13 /// Write all bytes in `buf` to the sink.
14 ///
15 /// Returns `Err` if the write could not be completed.
16 fn write(&mut self, buf: &[u8]) -> Result<(), WriteError>;
17 /// Flush any buffered data to the underlying sink.
18 fn flush(&mut self) -> Result<(), WriteError>;
19}
20
21#[cfg(not(feature = "std"))]
22impl Write for Vec<u8> {
23 fn write(&mut self, buf: &[u8]) -> Result<(), WriteError> {
24 self.extend_from_slice(buf);
25 Ok(())
26 }
27 fn flush(&mut self) -> Result<(), WriteError> {
28 Ok(())
29 }
30}
31
32#[cfg(feature = "std")]
33impl<W: std::io::Write> Write for W {
34 fn write(&mut self, buf: &[u8]) -> Result<(), WriteError> {
35 self.write_all(buf)?;
36 Ok(())
37 }
38 fn flush(&mut self) -> Result<(), WriteError> {
39 self.flush()?;
40 Ok(())
41 }
42}
43
44/// Writes CSV data to a [`Write`] sink.
45///
46/// Fields containing the delimiter, a newline, or a double-quote are
47/// automatically quoted. `""` escape sequences are used for quotes
48/// within quoted fields. Empty fields are written unquoted (as nothing),
49/// which is what every mainstream CSV writer does.
50///
51/// Internal buffering is used to avoid many small writes. Callers should
52/// not wrap the writer in a `BufWriter`.
53///
54/// # Security
55///
56/// This writer does **not** defend against CSV / formula injection
57/// ([CWE-1236]). Quoting a field does not neutralise it: spreadsheet
58/// applications still evaluate a quoted cell whose first character is
59/// `=`, `+`, `-`, `@`, tab or carriage return. If the data being written is
60/// attacker-controlled and the output is opened in a spreadsheet, sanitise
61/// those leading characters yourself (for example by prefixing an
62/// apostrophe) or disable formula evaluation in the consuming application.
63///
64/// [CWE-1236]: https://cwe.mitre.org/data/definitions/1236.html
65///
66/// # Example
67///
68/// ```no_run
69/// use csv::Writer;
70///
71/// let mut w = Writer::new(Vec::new());
72/// w.write_row(["name", "age", "city"])?;
73/// w.write_row(["Alice", "30", "New York, NY"])?;
74/// let result = String::from_utf8(w.into_inner()?).unwrap();
75/// assert_eq!(result, "name,age,city\r\nAlice,30,\"New York, NY\"\r\n");
76/// # Ok::<_, csv::WriteError>(())
77/// ```
78pub struct Writer<W: Write> {
79 pub(crate) writer: Option<W>,
80 pub(crate) delimiter: u8,
81 pub(crate) flexible: bool,
82 pub(crate) num_fields: Option<usize>,
83 pub(crate) row_count: usize,
84 pub(crate) buf: Vec<u8>,
85 pub(crate) headers: Option<Vec<String>>,
86 pub(crate) headers_written: bool,
87 /// If true, emit a UTF-8 byte-order mark before any other output.
88 pub(crate) write_bom: bool,
89 /// True once the byte-order mark has been emitted (or when none is wanted).
90 pub(crate) bom_written: bool,
91 #[cfg(feature = "serde")]
92 pub(crate) ser_values: Vec<Option<Vec<u8>>>,
93}
94
95impl<W: Write> Writer<W> {
96 /// Create a new writer wrapping the given output sink.
97 ///
98 /// The writer starts with default comma delimiter and strict
99 /// field-count validation (all rows must have the same number of fields).
100 pub fn new(writer: W) -> Self {
101 Writer {
102 writer: Some(writer),
103 delimiter: b',',
104 flexible: false,
105 num_fields: None,
106 row_count: 0,
107 buf: Vec::with_capacity(8192),
108 headers: None,
109 headers_written: false,
110 write_bom: false,
111 bom_written: false,
112 #[cfg(feature = "serde")]
113 ser_values: Vec::new(),
114 }
115 }
116
117 /// Set the field delimiter byte (default is `,`).
118 ///
119 /// # Panics
120 ///
121 /// Panics if `byte` is `"`, `\r` or `\n`; see
122 /// [`Reader::set_delimiter`](crate::Reader::set_delimiter).
123 ///
124 /// ```no_run
125 /// use csv::Writer;
126 /// let mut w = Writer::new(Vec::new()).set_delimiter(b'\t');
127 /// w.write_row(["a", "b", "c"])?;
128 /// # Ok::<_, csv::WriteError>(())
129 /// ```
130 pub fn set_delimiter(mut self, byte: u8) -> Self {
131 assert!(
132 byte != b'"' && byte != b'\r' && byte != b'\n',
133 "invalid CSV delimiter {byte:#04x}: `\"`, `\\r` and `\\n` are not allowed"
134 );
135 self.delimiter = byte;
136 self
137 }
138
139 /// Set whether variable field counts are allowed (default is `false`).
140 ///
141 /// When `false` (strict), all rows must have the same number of fields.
142 /// When `true`, rows may vary in field count.
143 pub fn set_flexible(mut self, yes: bool) -> Self {
144 self.flexible = yes;
145 self
146 }
147
148 /// Emits a UTF-8 byte-order mark (`U+FEFF`, the bytes `EF BB BF`) before any
149 /// other output (default is `false`).
150 ///
151 /// Enable this when the consumer expects one — Excel writes UTF-8 CSV with a
152 /// BOM, and some tools refuse to detect the encoding without it. The mark is
153 /// written once, before the first row, header or serialized record, and
154 /// [`Reader`](crate::Reader) strips a leading BOM by default, so output
155 /// round-trips.
156 pub fn set_write_bom(mut self, yes: bool) -> Self {
157 self.write_bom = yes;
158 self
159 }
160
161 /// Emits the byte-order mark into the buffer if it is wanted and has not
162 /// been written yet.
163 ///
164 /// Called at the top of every write entry point so the mark always precedes
165 /// the first piece of output.
166 pub(crate) fn write_pending_bom(&mut self) {
167 if self.write_bom && !self.bom_written {
168 self.buf.extend_from_slice(&crate::BOM);
169 self.bom_written = true;
170 }
171 }
172
173 /// Store column names for schema awareness.
174 ///
175 /// This sets the expected header names for future serde support and strict
176 /// field-count validation. It does **not** write anything to the output.
177 /// Use [`write_headers`](Self::write_headers) to emit the header row.
178 pub fn set_headers(mut self, headers: Vec<String>) -> Self {
179 self.headers = Some(headers);
180 self
181 }
182
183 /// Returns the stored header names, if any.
184 pub fn headers(&self) -> Option<&[String]> {
185 self.headers.as_deref()
186 }
187
188 /// Write a header row to the output and store the column names internally.
189 ///
190 /// The header row is written as a single CSV line. The field count from
191 /// this row becomes the expected count for all subsequent rows (unless
192 /// flexible mode is enabled).
193 ///
194 /// If this method has already been called, subsequent calls return
195 /// [`WriteError::HeadersAlreadyWritten`].
196 ///
197 /// # Errors
198 ///
199 /// Returns [`WriteError::HeadersAlreadyWritten`] if headers have already
200 /// been written. Returns [`WriteError::InconsistentFieldCount`] if the
201 /// number of headers differs from a prior [`write_row`](Self::write_row)
202 /// or from headers stored via [`set_headers`](Self::set_headers). The
203 /// count is validated before anything is written, so a rejected call
204 /// leaves the output untouched. Returns [`WriteError::Io`] if the
205 /// underlying writer fails.
206 pub fn write_headers<I, T>(&mut self, headers: I) -> Result<(), WriteError>
207 where
208 I: IntoIterator<Item = T>,
209 T: AsRef<str>,
210 {
211 if self.headers_written {
212 return Err(WriteError::HeadersAlreadyWritten);
213 }
214
215 let strings: Vec<String> = headers.into_iter().map(|s| s.as_ref().to_string()).collect();
216 let count = strings.len();
217
218 // Validate before writing anything so a mismatch cannot leave a
219 // partial or out-of-order row in the output.
220 if !self.flexible {
221 let expected = self.num_fields.or_else(|| self.headers.as_ref().map(|h| h.len()));
222 if let Some(expected) = expected
223 && count != expected
224 {
225 return Err(WriteError::InconsistentFieldCount {
226 expected,
227 found: count,
228 row: self.row_count + 1,
229 });
230 }
231 }
232
233 self.write_pending_bom();
234
235 for (i, s) in strings.iter().enumerate() {
236 if i > 0 {
237 self.buf.push(self.delimiter);
238 }
239 self.write_field(s.as_bytes())?;
240 }
241
242 self.buf.extend_from_slice(b"\r\n");
243 self.row_count += 1;
244 self.headers = Some(strings);
245 self.headers_written = true;
246 self.num_fields = Some(count);
247
248 if self.buf.len() >= 8192 {
249 self.flush()?;
250 }
251
252 Ok(())
253 }
254
255 /// Write a single row.
256 ///
257 /// Each element in the iterator is written as a CSV field. Fields are
258 /// auto-quoted if they contain the delimiter, a newline (`\n` or `\r`),
259 /// or a double-quote (`"`). Empty fields are written unquoted.
260 ///
261 /// # Round-tripping
262 ///
263 /// A row consisting of a single empty field is written as a blank line,
264 /// which readers (this one included, like most CSV readers) treat as an
265 /// empty line and skip. Such a row therefore cannot be read back. Every
266 /// other row round-trips exactly.
267 ///
268 /// # Errors
269 ///
270 /// Returns [`WriteError::InconsistentFieldCount`] if the number of
271 /// fields differs from previous rows (unless flexible mode is enabled).
272 /// Returns [`WriteError::Io`] if the underlying writer fails.
273 pub fn write_row<I, T>(&mut self, row: I) -> Result<(), WriteError>
274 where
275 I: IntoIterator<Item = T>,
276 T: AsRef<[u8]>,
277 {
278 self.write_pending_bom();
279
280 if self.flexible {
281 let mut first = true;
282 let mut field_count = 0;
283 for field in row {
284 if !first {
285 self.buf.push(self.delimiter);
286 }
287 first = false;
288 self.write_field(field.as_ref())?;
289 field_count += 1;
290 }
291 if self.num_fields.is_none() && field_count > 0 {
292 self.num_fields = Some(field_count);
293 }
294 self.buf.extend_from_slice(b"\r\n");
295 self.row_count += 1;
296 } else {
297 let buf_start = self.buf.len();
298 let mut field_count = 0;
299 for (i, field) in row.into_iter().enumerate() {
300 if i > 0 {
301 self.buf.push(self.delimiter);
302 }
303 self.write_field(field.as_ref())?;
304 field_count += 1;
305 }
306
307 match self.num_fields {
308 Some(expected) if field_count != expected => {
309 self.buf.truncate(buf_start);
310 return Err(WriteError::InconsistentFieldCount {
311 expected,
312 found: field_count,
313 row: self.row_count + 1,
314 });
315 }
316 None => {
317 self.num_fields = Some(field_count);
318 }
319 _ => {}
320 }
321
322 self.buf.extend_from_slice(b"\r\n");
323 self.row_count += 1;
324 }
325
326 if self.buf.len() >= 8192 {
327 self.flush()?;
328 }
329
330 Ok(())
331 }
332
333 fn write_field(&mut self, field: &[u8]) -> Result<(), WriteError> {
334 write_csv_field(&mut self.buf, self.delimiter, field);
335 Ok(())
336 }
337
338 /// Flush the internal buffer to the underlying writer.
339 pub fn flush(&mut self) -> Result<(), WriteError> {
340 if let Some(writer) = self.writer.as_mut() {
341 Write::write(writer, &self.buf)?;
342 Write::flush(writer)?;
343 }
344 self.buf.clear();
345 Ok(())
346 }
347
348 /// Flush any remaining data and return
349 /// the underlying writer.
350 pub fn into_inner(mut self) -> Result<W, WriteError> {
351 self.flush()?;
352 Ok(self.writer.take().unwrap())
353 }
354}
355
356impl<W: Write> Drop for Writer<W> {
357 fn drop(&mut self) {
358 if self.writer.is_some() {
359 let _ = self.flush();
360 }
361 }
362}
363
364/// Write a single CSV field to `buf`, applying quoting as needed.
365///
366/// Writes directly into `buf` to avoid intermediate buffering.
367///
368/// Empty fields are written as nothing rather than as `""`: both forms parse
369/// back to an empty field, but the unquoted form matches every mainstream CSV
370/// writer and avoids roughly tripling the output size of sparse data.
371pub(crate) fn write_csv_field(buf: &mut Vec<u8>, delimiter: u8, field: &[u8]) {
372 let needs_quoting =
373 memchr::memchr3(delimiter, b'\r', b'\n', field).is_some() || memchr::memchr(b'"', field).is_some();
374
375 if !needs_quoting {
376 buf.extend_from_slice(field);
377 return;
378 }
379
380 buf.push(b'"');
381
382 let mut pos = 0;
383 while let Some(quote_offset) = memchr::memchr(b'"', &field[pos..]) {
384 let abs_pos = pos + quote_offset;
385 buf.extend_from_slice(&field[pos..abs_pos]);
386 buf.extend_from_slice(b"\"\"");
387 pos = abs_pos + 1;
388 }
389 buf.extend_from_slice(&field[pos..]);
390 buf.push(b'"');
391}