diff_priv/test/
csv_exporter.rs

1use crate::test::environment::Datasets;
2use csv::Writer;
3
4const MUELLER_HEADERS: [&str; 16] = [
5    "uuid",
6    "disclosure_risk",
7    "id",
8    "gender",
9    "age",
10    "run_id",
11    "running",
12    "speed",
13    "duration",
14    "distance",
15    "side",
16    "bout",
17    "freq",
18    "walk_ratio",
19    "start",
20    "end",
21];
22
23const ADULT_LARGE_HEADERS: [&str; 15] = [
24    "uuid",
25    "disclosure_risk",
26    "timestamp",
27    "age",
28    "fnlwgt",
29    "education_num",
30    "capital_gain",
31    "capital_loss",
32    "hours_per_week",
33    "education",
34    "marital_status",
35    "workclass",
36    "native_country",
37    "occupation",
38    "class",
39];
40
41const ADULT_HEADERS: [&str; 10] = [
42    "uuid",
43    "disclosure_risk",
44    "timestamp",
45    "age",
46    "fnlwgt",
47    "education_num",
48    "capital_gain",
49    "capital_loss",
50    "hours_per_week",
51    "class",
52];
53
54pub struct CsvExporter {
55    data: Vec<Vec<String>>,
56    pub path: String,
57    dataset: Datasets,
58}
59
60impl CsvExporter {
61    pub fn new(path: String, dataset: Datasets) -> Self {
62        Self {
63            path,
64            data: vec![],
65            dataset,
66        }
67    }
68
69    pub fn export(&self) -> Result<(), Box<dyn std::error::Error>> {
70        println!("Exporting to {}", self.path);
71        let mut writer = Writer::from_path(&self.path)?;
72        // write header
73        match self.dataset {
74            Datasets::Adult(_) => writer.write_record(&ADULT_HEADERS)?,
75            Datasets::AdultLarge(_) => writer.write_record(&ADULT_LARGE_HEADERS)?,
76            Datasets::Mueller(_) => writer.write_record(&MUELLER_HEADERS)?,
77        }
78
79        for record in &self.data {
80            writer.write_record(record)?
81        }
82
83        writer.flush()?;
84        Ok(())
85    }
86
87    pub fn add(&mut self, value: Vec<String>) {
88        self.data.push(value)
89    }
90}
diff_priv/test/csv_exporter.rs

diff_priv/test/
csv_exporter.rs
