use std::fs::File;
use std::io::{BufReader, BufWriter, Write};
use std::path::{Path, PathBuf};
use tempfile::NamedTempFile;
use crate::capacity::{generate_alpha_columns, is_likely_data_row};
use crate::error::{SqawkError, SqawkResult};
use crate::table::{Table, Value};
pub struct DelimHandler {}
impl Default for DelimHandler {
fn default() -> Self {
Self::new()
}
}
impl DelimHandler {
pub fn new() -> Self {
DelimHandler {}
}
pub fn save_delimited(
&self,
table: &Table,
file_path: &Path,
delimiter: &str,
) -> SqawkResult<()> {
let parent = file_path.parent().unwrap_or(Path::new("."));
let temp_file = NamedTempFile::new_in(parent).map_err(SqawkError::IoError)?;
{
let mut writer = BufWriter::new(temp_file.as_file());
let header = table.columns().join(delimiter);
writeln!(writer, "{}", header).map_err(SqawkError::IoError)?;
for row in table.rows() {
let row_values: Vec<String> = row.iter().map(|value| value.to_string()).collect();
let row_str = row_values.join(delimiter);
writeln!(writer, "{}", row_str).map_err(SqawkError::IoError)?;
}
writer.flush().map_err(SqawkError::IoError)?;
}
temp_file
.as_file()
.sync_all()
.map_err(SqawkError::IoError)?;
temp_file
.persist(file_path)
.map_err(|e| SqawkError::IoError(e.error))?;
Ok(())
}
pub fn load_delimited(
&self,
file_spec: &str,
delimiter: &str,
custom_columns: Option<Vec<String>>,
) -> SqawkResult<Table> {
let (table_name, file_path) = self.parse_file_spec(file_spec)?;
let file = File::open(&file_path)?;
let reader = BufReader::new(file);
let delimiter_byte = if delimiter.len() == 1 {
delimiter.as_bytes()[0]
} else if delimiter == "\\t" {
b'\t' } else {
return Err(SqawkError::InvalidFileSpec(format!(
"Invalid delimiter: {}. Must be a single character.",
delimiter
)));
};
let has_headers = true;
let mut csv_reader = csv::ReaderBuilder::new()
.has_headers(has_headers)
.delimiter(delimiter_byte)
.comment(Some(b'#')) .flexible(true) .from_reader(reader);
let headers = if let Some(columns) = custom_columns {
columns
} else {
match csv_reader.headers().map_err(SqawkError::CsvError) {
Ok(header_row) => {
let fields: Vec<&str> = header_row.iter().collect();
if is_likely_data_row(&fields) {
generate_alpha_columns(fields.len())
} else {
fields.iter().map(|s| s.to_string()).collect()
}
}
Err(_) => {
let record_iter = csv_reader.records();
let first_record = record_iter.into_iter().next();
if let Some(Ok(record)) = first_record {
generate_alpha_columns(record.len())
} else {
vec!["a".to_string()]
}
}
}
};
let mut table = Table::new_with_delimiter(
&table_name,
headers,
Some(file_path.clone()),
delimiter.to_string(),
);
let column_count = table.column_count();
let mut row_buffer = Vec::with_capacity(column_count);
for result in csv_reader.records() {
let record = result.map_err(SqawkError::CsvError)?;
row_buffer.clear();
for field in record.iter() {
row_buffer.push(Value::from(field));
}
table.add_row_from_slice(&row_buffer)?;
}
Ok(table)
}
fn parse_file_spec(&self, file_spec: &str) -> SqawkResult<(String, PathBuf)> {
if let Some((table_name, file_path)) = file_spec.split_once('=') {
Ok((table_name.to_string(), PathBuf::from(file_path)))
} else {
let path = PathBuf::from(file_spec);
path.file_name().ok_or_else(|| {
SqawkError::InvalidFileSpec(format!("Invalid file specification: {}", file_spec))
})?;
let stem = path.file_stem().ok_or_else(|| {
SqawkError::InvalidFileSpec(format!("Invalid file specification: {}", file_spec))
})?;
Ok((stem.to_string_lossy().to_string(), path))
}
}
}