use crate::error::{Error, Result};
use serde::{Deserialize, Serialize};
use std::path::Path;
#[derive(Debug, Clone)]
pub enum DataSource {
Table(Table),
Series(Vec<Series>),
}
impl DataSource {
pub fn from_csv(path: impl AsRef<Path>) -> Result<Self> {
let path = path.as_ref();
let mut reader = csv::Reader::from_path(path)?;
let headers: Vec<String> = reader.headers()?.iter().map(String::from).collect();
let mut rows = Vec::new();
for result in reader.records() {
let record = result?;
let row: Vec<Value> = record.iter().map(Value::parse).collect();
rows.push(row);
}
Ok(Self::Table(Table { headers, rows }))
}
pub fn from_csv_string(content: &str) -> Result<Self> {
let mut reader = csv::Reader::from_reader(content.as_bytes());
let headers: Vec<String> = reader.headers()?.iter().map(String::from).collect();
let mut rows = Vec::new();
for result in reader.records() {
let record = result?;
let row: Vec<Value> = record.iter().map(Value::parse).collect();
rows.push(row);
}
Ok(Self::Table(Table { headers, rows }))
}
pub fn from_json(content: &str) -> Result<Self> {
let value: serde_json::Value = serde_json::from_str(content)?;
Self::from_json_value(value)
}
pub fn from_json_value(value: serde_json::Value) -> Result<Self> {
match value {
serde_json::Value::Array(arr) => {
if arr.is_empty() {
return Err(Error::NoData);
}
if arr[0].is_object() {
let headers: Vec<String> =
arr[0].as_object().unwrap().keys().cloned().collect();
let rows: Vec<Vec<Value>> = arr
.iter()
.filter_map(|obj| {
obj.as_object().map(|o| {
headers
.iter()
.map(|h| Value::from_json(o.get(h).cloned()))
.collect()
})
})
.collect();
Ok(Self::Table(Table { headers, rows }))
} else {
let values: Vec<f64> = arr.iter().filter_map(|v| v.as_f64()).collect();
Ok(Self::Series(vec![Series {
name: "data".to_string(),
values,
}]))
}
}
_ => Err(Error::InvalidData {
message: "expected JSON array".to_string(),
}),
}
}
#[must_use]
pub fn from_points(points: Vec<(f64, f64)>) -> Self {
let headers = vec!["x".to_string(), "y".to_string()];
let rows: Vec<Vec<Value>> = points
.into_iter()
.map(|(x, y)| vec![Value::Number(x), Value::Number(y)])
.collect();
Self::Table(Table { headers, rows })
}
#[must_use]
pub fn as_table(&self) -> Option<&Table> {
match self {
Self::Table(t) => Some(t),
_ => None,
}
}
#[must_use]
pub fn as_series(&self) -> Option<&[Series]> {
match self {
Self::Series(s) => Some(s),
_ => None,
}
}
}
#[derive(Debug, Clone)]
pub struct Table {
pub headers: Vec<String>,
pub rows: Vec<Vec<Value>>,
}
impl Table {
#[must_use]
pub fn column_index(&self, name: &str) -> Option<usize> {
self.headers.iter().position(|h| h == name)
}
#[must_use]
pub fn column_as_f64(&self, name: &str) -> Option<Vec<f64>> {
let idx = self.column_index(name)?;
Some(
self.rows
.iter()
.filter_map(|row| row.get(idx).and_then(Value::as_f64))
.collect(),
)
}
#[must_use]
pub fn column_as_str(&self, name: &str) -> Option<Vec<String>> {
let idx = self.column_index(name)?;
Some(
self.rows
.iter()
.filter_map(|row| row.get(idx).map(Value::to_string))
.collect(),
)
}
#[must_use]
pub fn row_count(&self) -> usize {
self.rows.len()
}
#[must_use]
pub fn column_count(&self) -> usize {
self.headers.len()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Series {
pub name: String,
pub values: Vec<f64>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Value {
Number(f64),
String(String),
Null,
}
impl Value {
#[must_use]
pub fn parse(s: &str) -> Self {
if s.is_empty() {
Self::Null
} else if let Ok(n) = s.parse::<f64>() {
Self::Number(n)
} else {
Self::String(s.to_string())
}
}
#[must_use]
pub fn from_json(value: Option<serde_json::Value>) -> Self {
match value {
Some(serde_json::Value::Number(n)) => Self::Number(n.as_f64().unwrap_or(0.0)),
Some(serde_json::Value::String(s)) => Self::String(s),
Some(serde_json::Value::Null) | None => Self::Null,
Some(v) => Self::String(v.to_string()),
}
}
#[must_use]
pub fn as_f64(&self) -> Option<f64> {
match self {
Self::Number(n) => Some(*n),
Self::String(s) => s.parse().ok(),
Self::Null => None,
}
}
#[must_use]
pub fn is_null(&self) -> bool {
matches!(self, Self::Null)
}
}
impl std::fmt::Display for Value {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Number(n) => write!(f, "{n}"),
Self::String(s) => write!(f, "{s}"),
Self::Null => write!(f, ""),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_csv_string() {
let csv = "x,y\n1,10\n2,20\n3,30";
let data = DataSource::from_csv_string(csv).unwrap();
let table = data.as_table().unwrap();
assert_eq!(table.headers, vec!["x", "y"]);
assert_eq!(table.row_count(), 3);
}
#[test]
fn table_column_access() {
let csv = "name,value\nalpha,100\nbeta,200";
let data = DataSource::from_csv_string(csv).unwrap();
let table = data.as_table().unwrap();
assert_eq!(table.column_index("name"), Some(0));
assert_eq!(table.column_index("value"), Some(1));
assert_eq!(table.column_index("missing"), None);
let values = table.column_as_f64("value").unwrap();
assert_eq!(values, vec![100.0, 200.0]);
}
#[test]
fn from_json_array_of_objects() {
let json = r#"[
{"x": 1, "y": 10},
{"x": 2, "y": 20}
]"#;
let data = DataSource::from_json(json).unwrap();
let table = data.as_table().unwrap();
assert_eq!(table.row_count(), 2);
}
#[test]
fn from_points() {
let data = DataSource::from_points(vec![(1.0, 10.0), (2.0, 20.0)]);
let table = data.as_table().unwrap();
assert_eq!(table.headers, vec!["x", "y"]);
assert_eq!(table.row_count(), 2);
}
#[test]
fn value_parsing() {
assert!(matches!(Value::parse("42"), Value::Number(n) if (n - 42.0).abs() < f64::EPSILON));
assert!(matches!(Value::parse("hello"), Value::String(_)));
assert!(matches!(Value::parse(""), Value::Null));
}
#[test]
fn value_as_f64() {
assert_eq!(Value::Number(42.0).as_f64(), Some(42.0));
assert_eq!(Value::String("42".to_string()).as_f64(), Some(42.0));
assert_eq!(Value::String("hello".to_string()).as_f64(), None);
assert_eq!(Value::Null.as_f64(), None);
}
}