108 lines
2.8 KiB
Rust
108 lines
2.8 KiB
Rust
use anyhow::{Result, bail};
|
|
use std::io::{BufRead, Write};
|
|
|
|
pub struct CsvRecordReader<R> {
|
|
reader: R,
|
|
buffer: String,
|
|
}
|
|
|
|
impl<R: BufRead> CsvRecordReader<R> {
|
|
pub fn new(reader: R) -> Self {
|
|
Self {
|
|
reader,
|
|
buffer: String::new(),
|
|
}
|
|
}
|
|
|
|
pub fn next_record(&mut self) -> Result<Option<Vec<String>>> {
|
|
self.buffer.clear();
|
|
let mut line = String::new();
|
|
loop {
|
|
line.clear();
|
|
let bytes = self.reader.read_line(&mut line)?;
|
|
if bytes == 0 {
|
|
if self.buffer.is_empty() {
|
|
return Ok(None);
|
|
}
|
|
break;
|
|
}
|
|
self.buffer.push_str(&line);
|
|
if csv_record_complete(&self.buffer) {
|
|
break;
|
|
}
|
|
}
|
|
Ok(Some(parse_csv_record(
|
|
self.buffer.trim_end_matches(['\r', '\n']),
|
|
)?))
|
|
}
|
|
}
|
|
|
|
fn csv_record_complete(record: &str) -> bool {
|
|
let mut in_quotes = false;
|
|
let mut chars = record.chars().peekable();
|
|
while let Some(character) = chars.next() {
|
|
if character == '"' {
|
|
if in_quotes && chars.peek() == Some(&'"') {
|
|
let _ = chars.next();
|
|
} else {
|
|
in_quotes = !in_quotes;
|
|
}
|
|
}
|
|
}
|
|
!in_quotes
|
|
}
|
|
|
|
pub fn parse_csv_record(record: &str) -> Result<Vec<String>> {
|
|
let mut fields = Vec::new();
|
|
let mut current = String::new();
|
|
let mut chars = record.chars().peekable();
|
|
let mut in_quotes = false;
|
|
while let Some(character) = chars.next() {
|
|
match character {
|
|
'"' if in_quotes && chars.peek() == Some(&'"') => {
|
|
current.push('"');
|
|
let _ = chars.next();
|
|
}
|
|
'"' => in_quotes = !in_quotes,
|
|
',' if !in_quotes => {
|
|
fields.push(current);
|
|
current = String::new();
|
|
}
|
|
_ => current.push(character),
|
|
}
|
|
}
|
|
if in_quotes {
|
|
bail!("CSV record has an unterminated quote");
|
|
}
|
|
fields.push(current);
|
|
Ok(fields)
|
|
}
|
|
|
|
pub fn write_csv_record(writer: &mut impl Write, fields: &[String]) -> Result<()> {
|
|
for (index, field) in fields.iter().enumerate() {
|
|
if index > 0 {
|
|
writer.write_all(b",")?;
|
|
}
|
|
if field.contains([',', '"', '\n', '\r']) {
|
|
writer.write_all(b"\"")?;
|
|
writer.write_all(field.replace('"', "\"\"").as_bytes())?;
|
|
writer.write_all(b"\"")?;
|
|
} else {
|
|
writer.write_all(field.as_bytes())?;
|
|
}
|
|
}
|
|
writer.write_all(b"\n")?;
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn parses_quoted_fields() {
|
|
let row = parse_csv_record(r#"name,"hello, ""world""",42"#).unwrap();
|
|
assert_eq!(row, vec!["name", "hello, \"world\"", "42"]);
|
|
}
|
|
}
|