formats/
data_struct.rs

1use super::create_io_error;
2use super::serde_json;
3use super::*;
4use csv;
5use linked_hash_map::LinkedHashMap;
6use std::fmt;
7use std::io::Error;
8use yaml_rust::Yaml;
9
10#[derive(Debug, Clone, PartialEq)]
11pub enum Number {
12    Int(i64),
13    Float(f64),
14}
15
16impl fmt::Display for Number {
17    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
18        match *self {
19            Number::Int(v) => write!(f, "{}", v),
20            Number::Float(v) => write!(f, "{}", v),
21        }
22    }
23}
24
25#[derive(Debug, Clone)]
26pub enum BasicTypes {
27    String(String),
28    Number(Number),
29    Null,
30}
31
32impl ToString for BasicTypes {
33    fn to_string(&self) -> String {
34        match self {
35            BasicTypes::String(s) => s.to_owned(),
36            BasicTypes::Number(num) => format!("{}", num),
37            BasicTypes::Null => "null".to_string(),
38        }
39    }
40}
41
42#[derive(Debug, Clone)]
43pub struct Row {
44    pub values: Vec<BasicTypes>,
45}
46
47#[derive(Debug)]
48pub struct Tabular {
49    pub headers: Row,
50    pub data: Vec<Row>,
51}
52
53pub fn str_to_basictypes(v: String) -> BasicTypes {
54    match v.parse::<f64>() {
55        Ok(f_val) => {
56            if v.contains(".") {
57                return BasicTypes::Number(Number::Float(f_val));
58            } else {
59                let num = v.parse::<i64>().unwrap();
60                return BasicTypes::Number(Number::Int(num));
61            };
62        }
63        Err(_) => BasicTypes::String(v),
64    }
65}
66
67impl Row {
68    pub fn new(values: Vec<String>) -> Row {
69        let mut row = Row { values: vec![] };
70        for v in values {
71            row.values.push(str_to_basictypes(v));
72        }
73        row
74    }
75
76    pub fn from_iter<'a, T: Iterator<Item = &'a str>>(iter: T) -> Row {
77        Row::new(iter.map(|x| String::from(x)).collect())
78    }
79
80    pub fn as_vec(&self) -> &Vec<BasicTypes> {
81        &self.values
82    }
83
84    pub fn to_csv_vec(&self) -> Vec<String> {
85        self.values.iter().map(|v| v.to_string()).collect()
86    }
87
88    pub fn to_serde_map(&self, headers: &Row) -> serde_json::Map<String, serde_json::Value> {
89        let mut map = serde_json::Map::new();
90        for (i, v) in self.values.iter().enumerate() {
91            let serde_v = match v.clone() {
92                BasicTypes::String(s) => serde_json::Value::String(s),
93                BasicTypes::Number(n) => match n {
94                    Number::Int(n) => json!(n),
95                    Number::Float(n) => json!(n),
96                },
97                BasicTypes::Null => serde_json::Value::Null,
98            };
99            map.insert(headers.values[i].clone().to_string(), serde_v);
100        }
101        map
102    }
103
104    pub fn to_yaml_hash(&self, headers: &Row) -> Yaml {
105        let mut rv = LinkedHashMap::new();
106
107        for (i, v) in self.values.iter().enumerate() {
108            let yaml_val = match v.clone() {
109                BasicTypes::String(s) => Yaml::String(s),
110                BasicTypes::Number(n) => match n {
111                    Number::Int(n) => Yaml::Integer(n),
112                    Number::Float(n) => Yaml::Real(n.to_string()),
113                },
114                BasicTypes::Null => Yaml::Null,
115            };
116            rv.insert(
117                Yaml::String(headers.values[i].clone().to_string()),
118                yaml_val,
119            );
120        }
121        Yaml::Hash(rv)
122    }
123}
124
125impl Tabular {
126    pub fn new(headers: Row) -> Tabular {
127        Tabular {
128            headers: headers,
129            data: vec![],
130        }
131    }
132
133    pub fn add_row(&mut self, row: Row) {
134        self.data.push(row);
135    }
136
137    pub fn add_data_from_iter<T>(&mut self, iter: T)
138    where
139        T: Iterator<Item = Row>,
140    {
141        for x in iter {
142            self.data.push(x)
143        }
144    }
145
146    pub fn has_headers(&self) -> bool {
147        self.headers.values.len() > 0
148    }
149
150    pub fn has_data(&self) -> bool {
151        self.data.len() > 0
152    }
153
154    pub fn write_csv(&self, path: &str) -> Result<(), Error> {
155        let mut wtr = csv::Writer::from_path(path)?;
156        if self.has_headers() {
157            match wtr.write_record(self.headers.to_csv_vec()) {
158                Err(e) => {
159                    return Err(create_io_error(&format!("{:?}", e)));
160                }
161                _ => {}
162            }
163        }
164
165        for row in self.data.iter() {
166            match wtr.write_record(row.to_csv_vec()) {
167                Err(e) => {
168                    return Err(create_io_error(&format!("{:?}", e)));
169                }
170                _ => {}
171            }
172        }
173        Ok(())
174    }
175
176    pub fn get_output_headers_data(&self) -> Result<(&Row, &[Row]), Error> {
177        let headers;
178        let data;
179        if self.has_headers() {
180            headers = &self.headers;
181            data = self.data.as_slice();
182        } else {
183            if self.has_data() {
184                headers = &self.data[0];
185                data = &self.data[1..];
186            } else {
187                return Err(create_io_error("the tablular does not have data"));
188            }
189        }
190        Ok((&headers, data))
191    }
192
193    pub fn write_json(&self, path: &str) -> Result<(), Error> {
194        let (headers, data) = self.get_output_headers_data()?;
195        let data = data.iter().map(|row| row.to_serde_map(headers)).collect();
196        io_json::write_json_object(path, &data, has_env("PRETTY"))?;
197        Ok(())
198    }
199
200    pub fn write_yaml(&self, path: &str) -> Result<(), Error> {
201        let (headers, data) = self.get_output_headers_data()?;
202        let data = data.iter().map(|row| row.to_yaml_hash(headers)).collect();
203        let doc = Yaml::Array(data);
204        io_yaml::write_yaml_doc(path, &doc)?;
205        Ok(())
206    }
207}