use csv_core::{self, WriteResult, Writer as CoreWriter, WriterBuilder as CoreWriterBuilder};
use serde::Serialize;
use crate::error::{Error, ErrorKind, Result};
use crate::serializer::{serialize, serialize_header};
use crate::{QuoteStyle, Terminator};
#[derive(Debug)]
pub struct WriterBuilder {
builder: CoreWriterBuilder,
capacity: usize,
flexible: bool,
has_headers: bool,
}
impl Default for WriterBuilder {
fn default() -> WriterBuilder {
WriterBuilder {
builder: CoreWriterBuilder::default(),
capacity: 8 * (1 << 10),
flexible: false,
has_headers: true,
}
}
}
impl WriterBuilder {
pub fn build(&self) -> Writer {
Writer::new(self)
}
pub fn delimiter(&mut self, delimiter: u8) -> &mut WriterBuilder {
self.builder.delimiter(delimiter);
self
}
pub fn has_headers(&mut self, yes: bool) -> &mut WriterBuilder {
self.has_headers = yes;
self
}
pub fn flexible(&mut self, yes: bool) -> &mut WriterBuilder {
self.flexible = yes;
self
}
pub fn terminator(&mut self, term: Terminator) -> &mut WriterBuilder {
self.builder.terminator(term.to_core());
self
}
pub fn quote_style(&mut self, style: QuoteStyle) -> &mut WriterBuilder {
self.builder.quote_style(style.to_core());
self
}
pub fn quote(&mut self, quote: u8) -> &mut WriterBuilder {
self.builder.quote(quote);
self
}
pub fn double_quote(&mut self, yes: bool) -> &mut WriterBuilder {
self.builder.double_quote(yes);
self
}
pub fn escape(&mut self, escape: u8) -> &mut WriterBuilder {
self.builder.escape(escape);
self
}
pub fn build_iter<I: IntoIterator>(&self, iter: I) -> crate::Iter<I::IntoIter> {
crate::Iter::new(iter, self.build())
}
#[cfg(feature = "stream")]
pub fn build_stream<S>(&self, stream: S) -> crate::Stream<S> {
crate::Stream::new(stream, self.build())
}
}
#[derive(Debug)]
pub struct Writer {
core: CoreWriter,
state: WriterState,
}
#[derive(Debug)]
struct WriterState {
header: HeaderState,
flexible: bool,
first_field_count: Option<u64>,
fields_written: u64,
panicked: bool,
}
#[derive(Debug)]
enum HeaderState {
Write,
DidWrite,
DidNotWrite,
None,
}
impl Default for Writer {
fn default() -> Self {
WriterBuilder::default().build()
}
}
impl Writer {
fn new(builder: &WriterBuilder) -> Writer {
let header_state = if builder.has_headers {
HeaderState::Write
} else {
HeaderState::None
};
Writer {
core: builder.builder.build(),
state: WriterState {
header: header_state,
flexible: builder.flexible,
first_field_count: None,
fields_written: 0,
panicked: false,
},
}
}
pub fn serialize<S: Serialize>(&mut self, buf: &mut Vec<u8>, record: S) -> Result<()> {
if let HeaderState::Write = self.state.header {
let wrote_header = serialize_header(self, buf, &record)?;
if wrote_header {
self.write_terminator(buf)?;
self.state.header = HeaderState::DidWrite;
} else {
self.state.header = HeaderState::DidNotWrite;
};
}
serialize(self, buf, &record)?;
self.write_terminator(buf)?;
Ok(())
}
pub fn write_record<I, T>(&mut self, buf: &mut Vec<u8>, record: I) -> Result<()>
where
I: IntoIterator<Item = T>,
T: AsRef<[u8]>,
{
for field in record.into_iter() {
self.write_field_impl(buf, field)?;
}
self.write_terminator(buf)
}
pub fn write_field<T: AsRef<[u8]>>(&mut self, buf: &mut Vec<u8>, field: T) -> Result<()> {
self.write_field_impl(buf, field)
}
#[inline(always)]
fn write_field_impl<T: AsRef<[u8]>>(&mut self, buf: &mut Vec<u8>, field: T) -> Result<()> {
if self.state.fields_written > 0 {
self.write_delimiter(buf)?;
}
let field = field.as_ref();
extend(buf, 2 * field.len() + 2, |buf| {
let (res, nin, nout) = self.core.field(field, buf);
debug_assert_eq!(res, WriteResult::InputEmpty);
debug_assert_eq!(nin, field.len());
self.state.fields_written += 1;
nout
});
Ok(())
}
fn write_delimiter(&mut self, buf: &mut Vec<u8>) -> Result<()> {
extend(buf, 2, |buf| {
let (res, nout) = self.core.delimiter(buf);
debug_assert_eq!(res, WriteResult::InputEmpty);
nout
});
Ok(())
}
fn write_terminator(&mut self, buf: &mut Vec<u8>) -> Result<()> {
self.check_field_count()?;
extend(buf, 4, |buf| {
let (res, nout) = self.core.terminator(buf);
debug_assert_eq!(res, WriteResult::InputEmpty);
self.state.fields_written = 0;
nout
});
Ok(())
}
fn check_field_count(&mut self) -> Result<()> {
if !self.state.flexible {
match self.state.first_field_count {
None => {
self.state.first_field_count = Some(self.state.fields_written);
}
Some(expected) if expected != self.state.fields_written => {
return Err(Error::new(ErrorKind::UnequalLengths {
expected_len: expected,
len: self.state.fields_written,
}))
}
Some(_) => {}
}
}
Ok(())
}
}
fn extend(buf: &mut Vec<u8>, max: usize, f: impl FnOnce(&mut [u8]) -> usize) {
let len = buf.len();
buf.resize(len + max, 0);
let n = f(&mut buf[len..]);
buf.resize(len + n, 0);
}
#[cfg(test)]
mod tests {
use super::WriterBuilder;
use serde::{serde_if_integer128, Serialize};
fn buf_as_string(buf: Vec<u8>) -> String {
String::from_utf8(buf).unwrap()
}
#[test]
fn one_record() {
let mut wtr = WriterBuilder::default().build();
let mut buf = vec![];
wtr.write_record(&mut buf, &["a", "b", "c"]).unwrap();
assert_eq!(buf_as_string(buf), "a,b,c\n");
}
#[test]
fn one_empty_record() {
let mut wtr = WriterBuilder::default().build();
let mut buf = vec![];
wtr.write_record(&mut buf, &[""]).unwrap();
assert_eq!(buf_as_string(buf), "\"\"\n");
}
#[test]
fn two_empty_records() {
let mut wtr = WriterBuilder::default().build();
let mut buf = vec![];
wtr.write_record(&mut buf, &[""]).unwrap();
wtr.write_record(&mut buf, &[""]).unwrap();
assert_eq!(buf_as_string(buf), "\"\"\n\"\"\n");
}
#[test]
fn serialize_with_headers() {
#[derive(Serialize)]
struct Row {
foo: i32,
bar: f64,
baz: bool,
}
let mut wtr = WriterBuilder::default().build();
let mut buf = vec![];
wtr.serialize(
&mut buf,
Row {
foo: 42,
bar: 42.5,
baz: true,
},
)
.unwrap();
assert_eq!(buf_as_string(buf), "foo,bar,baz\n42,42.5,true\n");
}
#[test]
fn serialize_no_headers() {
#[derive(Serialize)]
struct Row {
foo: i32,
bar: f64,
baz: bool,
}
let mut wtr = WriterBuilder::default().has_headers(false).build();
let mut buf = vec![];
wtr.serialize(
&mut buf,
Row {
foo: 42,
bar: 42.5,
baz: true,
},
)
.unwrap();
assert_eq!(buf_as_string(buf), "42,42.5,true\n");
}
serde_if_integer128! {
#[test]
fn serialize_no_headers_128() {
#[derive(Serialize)]
struct Row {
foo: i128,
bar: f64,
baz: bool,
}
let mut wtr =
WriterBuilder::default().has_headers(false).build();
let mut buf = vec![];
wtr.serialize(&mut buf, Row {
foo: 9_223_372_036_854_775_808,
bar: 42.5,
baz: true,
}).unwrap();
assert_eq!(buf_as_string(buf), "9223372036854775808,42.5,true\n");
}
}
#[test]
fn serialize_tuple() {
let mut wtr = WriterBuilder::default().build();
let mut buf = vec![];
wtr.serialize(&mut buf, (true, 1.3, "hi")).unwrap();
assert_eq!(buf_as_string(buf), "true,1.3,hi\n");
}
}