vantage_dataset/mocks/
csv.rs

1//! CSV mock implementation for testing and examples
2
3use crate::traits::{DataSet, ReadableDataSet, ReadableValueSet, Result, ValueSet};
4use indexmap::IndexMap;
5use std::collections::HashMap;
6use vantage_core::util::error::{Context, vantage_error};
7use vantage_types::{Entity, Record, vantage_type_system};
8
9// CSV type system - everything is a string since CSV is text-based
10vantage_type_system! {
11    type_trait: CsvType,
12    method_name: csv_string,
13    value_type: String,
14    type_variants: [String]
15}
16
17// Implement String for CSV type system
18impl CsvType for String {
19    type Target = CsvTypeStringMarker;
20
21    fn to_csv_string(&self) -> String {
22        self.clone()
23    }
24
25    fn from_csv_string(value: String) -> Option<Self> {
26        Some(value)
27    }
28}
29
30// Implement i64 for CSV type system (stored as string)
31impl CsvType for i64 {
32    type Target = CsvTypeStringMarker;
33
34    fn to_csv_string(&self) -> String {
35        self.to_string()
36    }
37
38    fn from_csv_string(value: String) -> Option<Self> {
39        value.parse().ok()
40    }
41}
42
43// Variant detection for CSV (only strings)
44impl CsvTypeVariants {
45    pub fn from_csv_string(_value: &String) -> Option<Self> {
46        Some(CsvTypeVariants::String)
47    }
48}
49
50/// MockCsv contains hardcoded CSV data as strings
51#[derive(Debug, Clone)]
52pub struct MockCsv {
53    files: HashMap<String, String>,
54}
55
56impl Default for MockCsv {
57    fn default() -> Self {
58        Self::new()
59    }
60}
61
62impl MockCsv {
63    pub fn new() -> Self {
64        let mut files = HashMap::new();
65
66        // Add users.csv data - all strings
67        files.insert(
68            "users.csv".to_string(),
69            r#"id,name,email,age
701,Alice Johnson,alice@example.com,28
712,Bob Smith,bob@example.com,35
723,Charlie Brown,charlie@example.com,42
734,Diana Prince,diana@example.com,31"#
74                .to_string(),
75        );
76
77        // Add products.csv data - all strings
78        files.insert(
79            "products.csv".to_string(),
80            r#"id,name,price,category
81101,Laptop,999.99,Electronics
82102,Coffee Mug,12.50,Kitchen
83103,Notebook,5.99,Office
84104,Wireless Mouse,25.00,Electronics"#
85                .to_string(),
86        );
87
88        Self { files }
89    }
90
91    pub fn get_csv_file<T: Entity<AnyCsvType>>(&self, filename: &str) -> CsvFile<T> {
92        CsvFile::new(self.clone(), filename)
93    }
94
95    pub fn list_files(&self) -> impl Iterator<Item = &String> {
96        self.files.keys()
97    }
98
99    pub fn get_file_content(&self, filename: &str) -> Result<&str> {
100        self.files
101            .get(filename)
102            .map(|s| s.as_str())
103            .ok_or_else(|| vantage_error!("File {} not found", filename))
104    }
105}
106
107/// CsvFile represents a single CSV file that can be read as a dataset
108#[derive(Debug, Clone)]
109pub struct CsvFile<T: Entity<AnyCsvType>> {
110    csv_ds: MockCsv,
111    filename: String,
112    _phantom: std::marker::PhantomData<T>,
113}
114
115impl<T: Entity<AnyCsvType>> ValueSet for CsvFile<T> {
116    type Id = usize;
117    type Value = AnyCsvType; // CSV values are always strings
118}
119
120impl<T: Entity<AnyCsvType>> CsvFile<T> {
121    pub fn new(csv_ds: MockCsv, filename: &str) -> Self {
122        Self {
123            csv_ds,
124            filename: filename.to_string(),
125            _phantom: std::marker::PhantomData,
126        }
127    }
128}
129
130#[async_trait::async_trait]
131impl<T> DataSet<T> for CsvFile<T> where T: Entity<AnyCsvType> {}
132
133#[async_trait::async_trait]
134impl<T> ReadableDataSet<T> for CsvFile<T>
135where
136    T: Entity<AnyCsvType>,
137{
138    async fn list(&self) -> Result<IndexMap<Self::Id, T>> {
139        let values = self.list_values().await?;
140        let mut records = IndexMap::new();
141
142        for (id, record) in values {
143            let entity = T::try_from_record(&record)
144                .map_err(|_| vantage_error!("Failed to convert record to entity"))?;
145            records.insert(id, entity);
146        }
147
148        Ok(records)
149    }
150
151    async fn get(&self, id: impl Into<Self::Id> + Send) -> Result<Option<T>> {
152        let id = id.into();
153        let Some(record) = self.get_value(&id).await? else {
154            return Ok(None);
155        };
156        let entity = T::try_from_record(&record)
157            .map_err(|_| vantage_error!("Failed to convert record to entity"))?;
158        Ok(Some(entity))
159    }
160
161    async fn get_some(&self) -> Result<Option<(Self::Id, T)>> {
162        if let Some((id, record)) = self.get_some_value().await? {
163            let entity = T::try_from_record(&record)
164                .map_err(|_| vantage_error!("Failed to convert record to entity"))?;
165            Ok(Some((id, entity)))
166        } else {
167            Ok(None)
168        }
169    }
170}
171
172#[async_trait::async_trait]
173impl<T> ReadableValueSet for CsvFile<T>
174where
175    T: Entity<AnyCsvType>,
176{
177    async fn list_values(&self) -> Result<IndexMap<Self::Id, Record<Self::Value>>> {
178        let content = self
179            .csv_ds
180            .get_file_content(&self.filename)
181            .context("Failed to get CSV content")?;
182
183        let mut reader = csv::ReaderBuilder::new().from_reader(content.as_bytes());
184        let mut records = IndexMap::new();
185
186        let headers = reader
187            .headers()
188            .context("Failed to read CSV headers")?
189            .clone();
190
191        for (idx, result) in reader.records().enumerate() {
192            let csv_record = result.context("Failed to read CSV record")?;
193
194            // Convert CSV record to Record<AnyCsvType> (all CSV values are strings)
195            let mut csv_record_map = Record::new();
196            for (i, field) in csv_record.iter().enumerate() {
197                if let Some(header) = headers.get(i) {
198                    csv_record_map.insert(header.to_string(), AnyCsvType::new(field.to_string()));
199                }
200            }
201
202            records.insert(idx, csv_record_map);
203        }
204
205        Ok(records)
206    }
207
208    async fn get_value(&self, id: &Self::Id) -> Result<Option<Record<Self::Value>>> {
209        let content = self
210            .csv_ds
211            .get_file_content(&self.filename)
212            .context("Failed to get CSV content")?;
213
214        let mut reader = csv::ReaderBuilder::new().from_reader(content.as_bytes());
215
216        let headers = reader
217            .headers()
218            .context("Failed to read CSV headers")?
219            .clone();
220
221        for (idx, result) in reader.records().enumerate() {
222            if idx == *id {
223                let csv_record = result.context("Failed to read CSV record")?;
224
225                // Convert CSV record to Record<AnyCsvType>
226                let mut csv_record_map = Record::new();
227                for (i, field) in csv_record.iter().enumerate() {
228                    if let Some(header) = headers.get(i) {
229                        csv_record_map
230                            .insert(header.to_string(), AnyCsvType::new(field.to_string()));
231                    }
232                }
233
234                return Ok(Some(csv_record_map));
235            }
236        }
237
238        Ok(None)
239    }
240
241    async fn get_some_value(&self) -> Result<Option<(Self::Id, Record<Self::Value>)>> {
242        let values = self.list_values().await?;
243        Ok(values.into_iter().next())
244    }
245}
vantage_dataset/mocks/csv.rs

vantage_dataset/mocks/
csv.rs
