Skip to content

Commit

Permalink
Add CSV support
Browse files Browse the repository at this point in the history
Fixes #2
  • Loading branch information
lucidfrontier45 committed Sep 28, 2023
1 parent b42ccd2 commit d5b08fc
Show file tree
Hide file tree
Showing 6 changed files with 110 additions and 0 deletions.
4 changes: 4 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -8,3 +8,7 @@ edition = "2021"
[dependencies]
serde = { version = "1", features = ["derive"] }
serde_json = "1"
csv = { version = "1", optional = true }

[features]
csv = ["dep:csv"]
3 changes: 3 additions & 0 deletions src/backend.rs
Original file line number Diff line number Diff line change
@@ -1,2 +1,5 @@
pub mod json;
pub mod jsonlines;

#[cfg(feature = "csv")]
pub mod csv;
87 changes: 87 additions & 0 deletions src/backend/csv.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
use std::io::{Read, Write};

use serde::{de::DeserializeOwned, Serialize};

use crate::common::Result;

pub fn read<T: DeserializeOwned>(reader: impl Read) -> Result<Vec<T>> {
let mut rdr = csv::Reader::from_reader(reader);
let mut records: Vec<T> = Vec::new();
for result in rdr.deserialize() {
let record: T = result?;
records.push(record);
}
Ok(records)
}

pub fn write<T: Serialize>(writer: impl Write, records: &[T]) -> Result<()> {
let mut wtr = csv::Writer::from_writer(writer);
for record in records {
wtr.serialize(record)?;
}
wtr.flush()?;
Ok(())
}

#[cfg(test)]
mod test {
use std::io::Cursor;

use serde::{Deserialize, Serialize};

use super::{read, write};

#[derive(Debug, Deserialize, Serialize, PartialEq, Eq)]
struct Record {
id: u32,
name: String,
}

#[test]
fn test_read() {
let text = r#"
id,name
1,foo
2,bar"#;

let records: Vec<Record> = read(Cursor::new(text)).unwrap();
let expected = vec![
Record {
id: 1,
name: "foo".to_owned(),
},
Record {
id: 2,
name: "bar".to_owned(),
},
];
assert_eq!(expected, records);
}

#[test]
fn test_write() {
let records = vec![
Record {
id: 1,
name: "foo".to_owned(),
},
Record {
id: 2,
name: "bar".to_owned(),
},
];
let mut cursor = Cursor::new(vec![]);
write(&mut cursor, &records).unwrap();
let data = String::from_utf8(cursor.into_inner())
.unwrap()
.trim()
.to_string();
let expected = r#"
id,name
1,foo
2,bar"#
.trim()
.to_owned();
assert_eq!(expected, data);
}
}
12 changes: 12 additions & 0 deletions src/common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,10 @@ use std::{fmt::Display, path::Path};
pub enum FileFormat {
Json,
JsonLines,
#[cfg(feature = "csv")]
Csv,
#[cfg(feature = "yaml")]
Yaml,
}

impl TryFrom<&str> for FileFormat {
Expand All @@ -13,6 +17,10 @@ impl TryFrom<&str> for FileFormat {
match value.trim().to_lowercase().as_str() {
"json" => Ok(FileFormat::Json),
"jsonl" | "jsl" => Ok(FileFormat::JsonLines),
#[cfg(feature = "yaml")]
"csv" => Ok(FileFormat::Csv),
#[cfg(feature = "yaml")]
"yaml" | "yml" => Ok(FileFormat::Yaml),
_ => Err(format!("Unknown file format: {}", value)),
}
}
Expand All @@ -35,6 +43,10 @@ impl Display for FileFormat {
match self {
FileFormat::Json => write!(f, "json"),
FileFormat::JsonLines => write!(f, "jsonl"),
#[cfg(feature = "csv")]
FileFormat::Csv => write!(f, "csv"),
#[cfg(feature = "yaml")]
FileFormat::Yaml => write!(f, "yaml"),
}
}
}
Expand Down
2 changes: 2 additions & 0 deletions src/read.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@ pub fn read_records_from_reader<T: DeserializeOwned>(
match file_format {
FileFormat::Json => backend::json::read(reader),
FileFormat::JsonLines => backend::jsonlines::read(reader),
#[cfg(feature = "csv")]
FileFormat::Csv => backend::csv::read(reader),
}
}

Expand Down
2 changes: 2 additions & 0 deletions src/write.rs
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ pub fn write_records_to_writer<T: Serialize>(
match file_format {
FileFormat::Json => backend::json::write(writer, records),
FileFormat::JsonLines => backend::jsonlines::write(writer, records),
#[cfg(feature = "csv")]
FileFormat::Csv => backend::csv::write(writer, records),
}
}

Expand Down

0 comments on commit d5b08fc

Please sign in to comment.