use crate::csv::{CompressionMethod, CsvEncoder};
use crate::error::{ExcelError, Result};
use crate::fast_writer::StreamingZipWriter;
use crate::types::CellValue;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::Path;
pub struct CsvWriter {
zip_writer: Option<StreamingZipWriter<File>>,
direct_writer: Option<BufWriter<File>>,
row_count: u64,
buffer: Vec<u8>,
delimiter: u8,
quote_char: u8,
line_ending: &'static [u8],
}
impl CsvWriter {
pub fn new<P: AsRef<Path>>(path: P) -> Result<Self> {
let path_ref = path.as_ref();
let path_str = path_ref.to_str().unwrap_or("");
if path_str.ends_with(".csv.zst") || path_str.ends_with(".csv.zip") {
Self::with_compression(path_ref, CompressionMethod::Zstd, 3)
} else if path_str.ends_with(".csv.gz") {
Self::with_compression(path_ref, CompressionMethod::Deflate, 6)
} else {
let file = File::create(path_ref)
.map_err(|e| ExcelError::WriteError(format!("Failed to create CSV file: {}", e)))?;
Ok(CsvWriter {
zip_writer: None,
direct_writer: Some(BufWriter::new(file)),
row_count: 0,
buffer: Vec::with_capacity(4096),
delimiter: b',',
quote_char: b'"',
line_ending: b"\n",
})
}
}
pub fn with_compression<P: AsRef<Path>>(
path: P,
method: CompressionMethod,
level: u32,
) -> Result<Self> {
let path_ref = path.as_ref();
let mut zip = StreamingZipWriter::with_method(path_ref, method, level)
.map_err(|e| ExcelError::WriteError(format!("Failed to create ZIP writer: {}", e)))?;
let entry_name = path_ref
.file_stem()
.and_then(|s| s.to_str())
.map(|s| {
let clean = s
.trim_end_matches(".csv")
.trim_end_matches(".zst")
.trim_end_matches(".gz");
format!("{}.csv", clean)
})
.unwrap_or_else(|| "data.csv".to_string());
zip.start_entry(&entry_name)
.map_err(|e| ExcelError::WriteError(format!("Failed to start ZIP entry: {}", e)))?;
Ok(CsvWriter {
zip_writer: Some(zip),
direct_writer: None,
row_count: 0,
buffer: Vec::with_capacity(4096),
delimiter: b',',
quote_char: b'"',
line_ending: b"\n",
})
}
pub fn delimiter(mut self, delim: u8) -> Self {
self.delimiter = delim;
self
}
pub fn quote_char(mut self, quote: u8) -> Self {
self.quote_char = quote;
self
}
pub fn write_row<I, S>(&mut self, data: I) -> Result<()>
where
I: IntoIterator<Item = S>,
S: AsRef<str>,
{
self.buffer.clear();
let encoder = CsvEncoder::new(self.delimiter, self.quote_char);
let fields: Vec<String> = data.into_iter().map(|s| s.as_ref().to_string()).collect();
let refs: Vec<&str> = fields.iter().map(|s| s.as_str()).collect();
encoder.encode_row(&refs, &mut self.buffer);
self.buffer.extend_from_slice(self.line_ending);
if let Some(ref mut zip) = self.zip_writer {
zip.write_data(&self.buffer)
.map_err(|e| ExcelError::WriteError(format!("Failed to write to ZIP: {}", e)))?;
} else if let Some(ref mut writer) = self.direct_writer {
writer
.write_all(&self.buffer)
.map_err(|e| ExcelError::WriteError(format!("Failed to write to file: {}", e)))?;
}
self.row_count += 1;
Ok(())
}
pub fn write_row_typed(&mut self, cells: &[CellValue]) -> Result<()> {
let strings: Vec<String> = cells.iter().map(|c| c.as_string()).collect();
let refs: Vec<&str> = strings.iter().map(|s| s.as_str()).collect();
self.write_row(refs)
}
pub fn write_rows_batch<I, R, S>(&mut self, rows: I) -> Result<()>
where
I: IntoIterator<Item = R>,
R: IntoIterator<Item = S>,
S: AsRef<str>,
{
for row_data in rows {
self.write_row(row_data)?;
}
Ok(())
}
pub fn row_count(&self) -> u64 {
self.row_count
}
pub fn save(mut self) -> Result<()> {
if let Some(zip) = self.zip_writer.take() {
zip.finish()
.map_err(|e| ExcelError::WriteError(format!("Failed to finish ZIP: {}", e)))?;
} else if let Some(mut writer) = self.direct_writer.take() {
writer
.flush()
.map_err(|e| ExcelError::WriteError(format!("Failed to flush file: {}", e)))?;
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Read;
#[test]
fn test_plain_csv() -> Result<()> {
let path = "test_output.csv";
{
let mut writer = CsvWriter::new(path)?;
writer.write_row(["Name", "Age", "City"])?;
writer.write_row(["Alice", "30", "NYC"])?;
writer.save()?;
}
let mut content = String::new();
File::open(path)?.read_to_string(&mut content)?;
assert!(content.contains("Name,Age,City"));
assert!(content.contains("Alice,30,NYC"));
std::fs::remove_file(path).ok();
Ok(())
}
#[test]
fn test_typed_values() -> Result<()> {
let path = "test_typed.csv";
{
let mut writer = CsvWriter::new(path)?;
writer.write_row_typed(&[
CellValue::String("Test".to_string()),
CellValue::Int(42),
CellValue::Float(3.15),
])?;
writer.save()?;
}
let mut content = String::new();
File::open(path)?.read_to_string(&mut content)?;
assert!(content.contains("Test,42,3.15"));
std::fs::remove_file(path).ok();
Ok(())
}
#[test]
fn test_edge_cases() -> Result<()> {
let path = "test_edge.csv";
{
let mut writer = CsvWriter::new(path)?;
writer.write_row(["a,b", r#"Say "Hi""#, "Line1\nLine2"])?;
writer.save()?;
}
let mut content = String::new();
File::open(path)?.read_to_string(&mut content)?;
assert!(content.contains(r#""a,b""#));
assert!(content.contains(r#""Say ""Hi""""#));
std::fs::remove_file(path).ok();
Ok(())
}
}