use csv::WriterBuilder;
use serde::Serialize;
use std::error::Error;
use std::fs::{self, File};
use std::io::BufWriter;
use std::path::Path;
pub mod serializer;
pub struct DataStorage {
writer: Option<csv::Writer<BufWriter<File>>>,
cnt: usize,
}
impl DataStorage {
pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, Box<dyn Error>> {
let path = path.as_ref();
if let Some(parent) = path.parent().filter(|parent| !parent.exists()) {
fs::create_dir_all(parent)?;
}
let file = File::create(path)?;
let buf_writer = BufWriter::new(file);
Ok(Self {
writer: Some(Self::csv_writer(buf_writer, b',')),
cnt: 0,
})
}
fn csv_writer(buf_writer: BufWriter<File>, delimiter: u8) -> csv::Writer<BufWriter<File>> {
WriterBuilder::new()
.has_headers(false)
.delimiter(delimiter)
.from_writer(buf_writer)
}
pub fn set_separator(mut self, separator: char) -> Self {
if !separator.is_ascii() {
return self;
}
if let Some(writer) = self.writer.take() {
self.writer = Some(match writer.into_inner() {
Ok(buf_writer) => Self::csv_writer(buf_writer, separator as u8),
Err(err) => err.into_inner(),
});
}
self
}
pub fn set_header<I, T>(mut self, header: I) -> Self
where
I: IntoIterator<Item = T>,
T: AsRef<str>,
{
if let Some(writer) = &mut self.writer {
for name in header {
let _ = writer.write_field(name.as_ref());
}
let _ = writer.write_record(None::<&[u8]>);
}
self
}
pub fn add<T: Serialize>(&mut self, data: &T) -> Result<(), Box<dyn Error>> {
if let Some(writer) = &mut self.writer {
writer.serialize(data)?;
self.cnt += 1;
#[allow(clippy::manual_is_multiple_of)]
if self.cnt % 100 == 0 {
let _ = writer.flush();
}
}
Ok(())
}
pub fn flush(&mut self) -> Result<(), Box<dyn Error>> {
if let Some(writer) = &mut self.writer {
writer.flush()?;
}
Ok(())
}
pub fn close(mut self) -> Result<(), Box<dyn Error>> {
if let Some(mut writer) = self.writer.take() {
writer.flush()?;
}
Ok(())
}
}
impl Drop for DataStorage {
fn drop(&mut self) {
let _ = self.flush();
}
}