-
Notifications
You must be signed in to change notification settings - Fork 217
/
writer.rs
453 lines (422 loc) · 15.2 KB
/
writer.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
use std::borrow::IntoCow;
use std::error::FromError;
use std::fs;
use std::io::{self, Write};
use std::path::AsPath;
use std::str;
use rustc_serialize::Encodable;
use {
BorrowBytes, ByteString, Result, Encoded, Error, RecordTerminator,
};
/// The quoting style to use when writing CSV data.
#[derive(Copy)]
pub enum QuoteStyle {
/// This puts quotes around every field. Always.
Always,
/// This puts quotes around fields only when necessary.
///
/// They are necessary when fields are empty or contain a quote, delimiter
/// or record terminator.
///
/// This is the default.
Necessary,
/// This *never* writes quotes.
///
/// If a field requires quotes, then the writer will report an error.
Never,
}
/// A CSV writer.
///
/// This writer provides a convenient interface for encoding CSV data. While
/// creating CSV data is much easier than parsing it, having a writer can
/// be convenient because it can handle quoting for you automatically.
/// Moreover, this particular writer supports `Encodable` types, which makes
/// it easy to write your custom types as CSV records.
///
/// All CSV data produced by this writer, with default options, conforms with
/// [RFC 4180](https://tools.ietf.org/html/rfc4180). (If certain options like
/// flexible record lengths are enabled, then compliance with RFC 4180 cannot
/// be guaranteed.)
///
/// One slight deviation is that records with a single empty field are always
/// encoded as `""`. This ensures that the record is not skipped since some
/// CSV parsers will ignore consecutive record terminators (like the one in
/// this crate).
///
/// ### Example
///
/// Here's an example that encodes word pairs and their edit distances:
///
/// ```rust
/// let records = vec![
/// ("sticker", "mortals", 7u),
/// ("bribed", "personae", 7u),
/// ("wobbling", "poncing", 4u),
/// ("interposed", "emmett", 9u),
/// ("chocolate", "refile", 7u),
/// ];
///
/// let mut wtr = csv::Writer::from_memory();
/// for record in records.into_iter() {
/// let result = wtr.encode(record);
/// assert!(result.is_ok());
/// }
/// ```
pub struct Writer<W> {
buf: io::BufWriter<W>,
delimiter: u8,
record_terminator: RecordTerminator,
flexible: bool,
quote: u8,
escape: u8,
double_quote: bool,
quote_style: QuoteStyle,
first_len: usize,
}
impl Writer<fs::File> {
/// Creates a new `Writer` that writes CSV data to the file path given.
///
/// The file is created if it does not already exist and is truncated
/// otherwise.
pub fn from_file<P: AsPath + ?Sized>
(path: &P)
-> Result<Writer<fs::File>> {
Ok(Writer::from_writer(try!(fs::File::create(path))))
}
}
impl<W: io::Write> Writer<W> {
/// Creates a new CSV writer that writes to the `io::Write` given.
///
/// Note that the writer is buffered for you automatically.
pub fn from_writer(w: W) -> Writer<W> {
Writer::from_buffer(io::BufWriter::new(w))
}
/// Creates a new CSV writer that writes to the buffer given.
///
/// This lets you specify your own buffered writer (e.g., use a different
/// capacity). All other constructors wrap the writer given in a buffer
/// with default capacity.
pub fn from_buffer(buf: io::BufWriter<W>) -> Writer<W> {
Writer {
buf: buf,
delimiter: b',',
record_terminator: RecordTerminator::Any(b'\n'),
flexible: false,
quote: b'"',
escape: b'\\',
double_quote: true,
quote_style: QuoteStyle::Necessary,
first_len: 0,
}
}
}
impl Writer<Vec<u8>> {
/// Creates a new CSV writer that writes to an in memory buffer. At any
/// time, `to_string` or `to_bytes` can be called to retrieve the
/// cumulative CSV data.
pub fn from_memory() -> Writer<Vec<u8>> {
Writer::from_writer(Vec::with_capacity(1024 * 64))
}
/// Returns the written CSV data as a string.
pub fn as_string<'r>(&'r mut self) -> &'r str {
match self.buf.flush() {
// shouldn't panic with Vec<u8>
Err(err) => panic!("Error flushing to Vec<u8>: {}", err),
// This seems suspicious. If the client only writes `String`
// values, then this can never fail. If the client is writing
// byte strings, then they should be calling `to_bytes` instead.
Ok(()) => str::from_utf8(&**self.buf.get_ref()).unwrap(),
}
}
/// Returns the encoded CSV data as raw bytes.
pub fn as_bytes<'r>(&'r mut self) -> &'r [u8] {
match self.buf.flush() {
// shouldn't panic with Vec<u8>
Err(err) => panic!("Error flushing to Vec<u8>: {}", err),
Ok(()) => &**self.buf.get_ref(),
}
}
}
impl<W: io::Write> Writer<W> {
/// Writes a record by encoding any `Encodable` value.
///
/// This is the most convenient way to write CSV data. Most Rust types
/// map to CSV data in a straight forward way. A vector is just a sequence
/// of fields. Similarly for a struct. Enumerations of zero or one
/// arguments are supported too. (Enums with zero arguments encode to their
/// name, while enums of one argument encode to their constituent value.)
/// Option types are also supported (`None` encodes to an empty field).
///
/// ### Example
///
/// This example encodes word pairs that may or may not have their
/// edit distances computed.
///
/// ```rust
/// extern crate "rustc-serialize" as rustc_serialize;
/// # extern crate csv;
/// # fn main() {
///
/// #[derive(RustcEncodable)]
/// struct Distance {
/// name1: &'static str,
/// name2: &'static str,
/// dist: Option<usize>,
/// }
///
/// let records = vec![
/// Distance { name1: "sticker", name2: "mortals", dist: None },
/// Distance { name1: "bribed", name2: "personae", dist: Some(7) },
/// ];
///
/// let mut wtr = csv::Writer::from_memory();
/// for record in records.into_iter() {
/// let result = wtr.encode(record);
/// assert!(result.is_ok());
/// }
/// assert_eq!(wtr.as_string(),
/// "sticker,mortals,\nbribed,personae,7\n");
/// # }
/// ```
pub fn encode<E>(&mut self, e: E) -> Result<()> where E: Encodable {
let mut erecord = Encoded::new();
try!(e.encode(&mut erecord));
self.write(erecord.unwrap().into_iter())
}
/// Writes a record of strings (Unicode or raw bytes).
///
/// This is meant to be the standard method provided by most CSV writers.
/// That is, it writes a record of strings---no more and no less.
///
/// This method accepts an iterator of *fields* for a single record. Each
/// field must satisfy `BorrowBytes`, which allows the caller to control
/// allocation.
///
/// ### Example
///
/// This shows how to write string records.
///
/// ```rust
/// let records = vec![
/// vec!["sticker", "mortals", "7"],
/// vec!["bribed", "personae", "7"],
/// vec!["wobbling", "poncing", "4"],
/// vec!["interposed", "emmett", "9"],
/// vec!["chocolate", "refile", "7"],
/// ];
///
/// let mut wtr = csv::Writer::from_memory();
/// for record in records.into_iter() {
/// let result = wtr.write(record.into_iter());
/// assert!(result.is_ok());
/// }
/// ```
///
/// This shows how to write records that do not correspond to a valid UTF-8
/// encoding. (Note the use of Rust's byte string syntax!)
///
/// ```rust
/// let mut wtr = csv::Writer::from_memory();
/// let result = wtr.write(vec![b"\xff", b"\x00"].into_iter());
/// assert!(result.is_ok());
///
/// assert_eq!(wtr.as_bytes(), b"\xff,\x00\n");
/// ```
pub fn write<'a, I>(&mut self, r: I) -> Result<()>
where I: Iterator, <I as Iterator>::Item: BorrowBytes {
self.write_iter(r.map(|f| Ok(f)))
}
/// Writes a record of results. If any of the results resolve to an error,
/// then writing stops and that error is returned.
#[doc(hidden)]
pub fn write_iter<'a, I, F>(&mut self, r: I) -> Result<()>
where I: Iterator<Item=Result<F>>, F: BorrowBytes {
let delim = self.delimiter;
let mut count = 0;
let mut last_len = 0;
for field in r {
if count > 0 {
try!(self.w_bytes(&[delim]));
}
count += 1;
let field = try!(field);
last_len = field.borrow_bytes().len();
try!(self.w_user_bytes(field.borrow_bytes()));
}
// This tomfoolery makes sure that a record with a single empty field
// is encoded as `""`. Otherwise, you end up with a run of consecutive
// record terminators, which are ignored by some CSV parsers (such
// as the one in this library).
if count == 1 && last_len == 0 {
let q = self.quote;
try!(self.w_bytes(&[q, q]));
}
try!(self.w_lineterm());
self.set_first_len(count)
}
/// Flushes the underlying buffer.
pub fn flush(&mut self) -> Result<()> {
self.buf.flush().map_err(FromError::from_error)
}
}
impl<W: io::Write> Writer<W> {
/// The delimiter to use when writing CSV data.
///
/// Since the CSV writer is meant to be mostly encoding agnostic, you must
/// specify the delimiter as a single ASCII byte. For example, to write
/// tab-delimited data, you would use `b'\t'`.
///
/// The default value is `b','`.
pub fn delimiter(mut self, delimiter: u8) -> Writer<W> {
self.delimiter = delimiter;
self
}
/// Whether to allow flexible length records when writing CSV data.
///
/// When this is set to `true`, records in the CSV data can have different
/// lengths. By default, this is disabled, which will cause the CSV writer
/// to return an error if it tries to write a record that has a different
/// length than other records it has already written.
pub fn flexible(mut self, yes: bool) -> Writer<W> {
self.flexible = yes;
self
}
/// Sets the record terminator to use when writing CSV data.
///
/// By default, this is `RecordTerminator::Any(b'\n')`. If you want to
/// use CRLF (`\r\n`) line endings, then use `RecordTerminator:CRLF`.
pub fn record_terminator(mut self, term: RecordTerminator) -> Writer<W> {
self.record_terminator = term;
self
}
/// Set the quoting style to use when writing CSV data.
///
/// By default, this is set to `QuoteStyle::Necessary`, which will only
/// use quotes when they are necessary to preserve the integrity of data.
pub fn quote_style(mut self, style: QuoteStyle) -> Writer<W> {
self.quote_style = style;
self
}
/// Set the quote character to use when writing CSV data.
///
/// Since the CSV parser is meant to be mostly encoding agnostic, you must
/// specify the quote as a single ASCII byte. For example, to write
/// single quoted data, you would use `b'\''`.
///
/// The default value is `b'"'`.
pub fn quote(mut self, quote: u8) -> Writer<W> {
self.quote = quote;
self
}
/// Set the escape character to use when writing CSV data.
///
/// This is only used when `double_quote` is set to `false`.
///
/// Since the CSV parser is meant to be mostly encoding agnostic, you must
/// specify the escape as a single ASCII byte.
///
/// The default value is `b'\\'`.
pub fn escape(mut self, escape: u8) -> Writer<W> {
self.escape = escape;
self
}
/// Set the quoting escape mechanism.
///
/// When enabled (which is the default), quotes are escaped by doubling
/// them. e.g., `"` escapes to `""`.
///
/// When disabled, quotes are escaped with the escape character (which
/// is `\\` by default).
pub fn double_quote(mut self, yes: bool) -> Writer<W> {
self.double_quote = yes;
self
}
}
impl<W: io::Write> Writer<W> {
fn err<'a, S, T>(&self, msg: S) -> Result<T> where S: IntoCow<'a, str> {
Err(Error::Encode(msg.into_cow().into_owned()))
}
fn w_bytes(&mut self, s: &[u8]) -> Result<()> {
self.buf.write_all(s).map_err(Error::Io)
}
fn w_user_bytes(&mut self, s: &[u8]) -> Result<()> {
if try!(self.should_quote(s)) {
let quoted = self.quote_field(s);
self.w_bytes(&*quoted)
} else {
self.w_bytes(s)
}
}
fn w_lineterm(&mut self) -> Result<()> {
match self.record_terminator {
RecordTerminator::CRLF => self.w_bytes(b"\r\n"),
RecordTerminator::Any(b) => self.w_bytes(&[b]),
}
}
fn set_first_len(&mut self, cur_len: usize) -> Result<()> {
if cur_len == 0 {
return self.err("Records must have length greater than 0.")
}
if !self.flexible {
if self.first_len == 0 {
self.first_len = cur_len;
} else if self.first_len != cur_len {
return self.err(format!(
"Record has length {} but other records have length {}",
cur_len, self.first_len))
}
}
Ok(())
}
fn should_quote(&self, field: &[u8]) -> Result<bool> {
let needs = || field.iter().any(|&b| self.byte_needs_quotes(b));
match self.quote_style {
QuoteStyle::Always => Ok(true),
QuoteStyle::Necessary => Ok(needs()),
QuoteStyle::Never => {
if !needs() {
Ok(false)
} else {
self.err(format!(
"Field requires quotes, but quote style \
is 'Never': '{}'",
String::from_utf8_lossy(field)))
}
}
}
}
fn byte_needs_quotes(&self, b: u8) -> bool {
b == self.delimiter
|| self.record_terminator == b
|| b == self.quote
// This is a bit hokey. By default, the record terminator is
// '\n', but we still need to quote '\r' because the reader
// interprets '\r' as a record terminator by default.
|| b == b'\r' || b == b'\n'
}
fn quote_field(&self, mut s: &[u8]) -> ByteString {
let mut buf = Vec::with_capacity(s.len() + 2);
buf.push(self.quote);
loop {
match s.iter().position(|&v| v == self.quote) {
None => {
buf.extend(s.iter().map(|&x|x));
break
}
Some(next_quote) => {
buf.extend(s[..next_quote].iter().map(|&x|x));
if self.double_quote {
buf.push(self.quote);
buf.push(self.quote);
} else {
buf.push(self.escape);
buf.push(self.quote);
}
s = &s[next_quote + 1..];
}
}
}
buf.push(self.quote);
ByteString::from_bytes(buf)
}
}