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}