prettytable/
csv.rs

1//! CSV impl and reexported types
2
3extern crate csv;
4
5pub use self::csv::{Reader, Writer, Result, ReaderBuilder};
6use crate::AsTableSlice;
7use std::path::Path;
8use std::io::{Read, Write};
9
10impl<'a> super::TableSlice<'a> {
11    /// Write the table to the specified writer.
12    pub fn to_csv<W: Write>(&self, w: W) -> Result<Writer<W>> {
13        self.to_csv_writer(Writer::from_writer(w))
14    }
15
16    /// Write the table to the specified writer.
17    ///
18    /// This allows for format customisation.
19    pub fn to_csv_writer<W: Write>(&self,
20                                mut writer: Writer<W>)
21                                -> Result<Writer<W>> {
22        for title in self.titles {
23            writer.write_record(title.iter().map(|c| c.get_content()))?;
24        }
25        for row in self.rows {
26            writer.write_record(row.iter().map(|c| c.get_content()))?;
27        }
28
29        writer.flush()?;
30        Ok(writer)
31    }
32}
33
34impl super::Table {
35    /// Create a table from a CSV string
36    ///
37    /// For more customisability use `from_csv()`
38    pub fn from_csv_string(csv_s: &str) -> Result<Self> {
39        Ok(Self::from_csv(
40            &mut ReaderBuilder::new()
41                .has_headers(false)
42                .from_reader(csv_s.as_bytes())))
43    }
44
45    /// Create a table from a CSV file
46    ///
47    /// For more customisability use `from_csv()`
48    pub fn from_csv_file<P: AsRef<Path>>(csv_p: P) -> Result<Self> {
49        Ok(Self::from_csv(
50            &mut ReaderBuilder::new()
51                .has_headers(false)
52                .from_path(csv_p)?))
53    }
54
55    /// Create a table from a CSV reader
56    pub fn from_csv<R: Read>(reader: &mut Reader<R>) -> Self {
57        Self::init(reader
58                        .records()
59                        .map(|row| {
60                                super::Row::new(row.unwrap()
61                                            .into_iter()
62                                            .map(|cell| super::Cell::new(&cell))
63                                            .collect())
64                            })
65                        .collect())
66    }
67
68
69    /// Write the table to the specified writer.
70    pub fn to_csv<W: Write>(&self, w: W) -> Result<Writer<W>> {
71        self.as_slice().to_csv(w)
72    }
73
74    /// Write the table to the specified writer.
75    ///
76    /// This allows for format customisation.
77    pub fn to_csv_writer<W: Write>(&self, writer: Writer<W>) -> Result<Writer<W>> {
78        self.as_slice().to_csv_writer(writer)
79    }
80}
81
82
83#[cfg(test)]
84mod tests {
85    use {Table, Row, Cell};
86
87    static CSV_S: &'static str = "ABC,DEFG,HIJKLMN\n\
88                                foobar,bar,foo\n\
89                                foobar2,bar2,foo2\n";
90
91    fn test_table() -> Table {
92        let mut table = Table::new();
93        table
94            .add_row(Row::new(vec![Cell::new("ABC"), Cell::new("DEFG"), Cell::new("HIJKLMN")]));
95        table.add_row(Row::new(vec![Cell::new("foobar"), Cell::new("bar"), Cell::new("foo")]));
96        table.add_row(Row::new(vec![Cell::new("foobar2"),
97                                    Cell::new("bar2"),
98                                    Cell::new("foo2")]));
99        table
100    }
101
102    #[test]
103    fn from() {
104        assert_eq!(test_table().to_string().replace("\r\n", "\n"),
105                    Table::from_csv_string(CSV_S)
106                        .unwrap()
107                        .to_string()
108                        .replace("\r\n", "\n"));
109    }
110
111    #[test]
112    fn to() {
113        assert_eq!(
114            String::from_utf8(
115                test_table()
116                    .to_csv(Vec::new())
117                    .unwrap()
118                    .into_inner()
119                    .unwrap()
120                ).unwrap(),
121                CSV_S);
122    }
123
124    #[test]
125    fn trans() {
126        assert_eq!(
127            Table::from_csv_string(
128                &String::from_utf8(
129                    test_table()
130                        .to_csv(Vec::new())
131                        .unwrap()
132                        .into_inner()
133                        .unwrap()
134                ).unwrap()
135            ).unwrap()
136            .to_string()
137            .replace("\r\n", "\n"),
138            test_table().to_string().replace("\r\n", "\n"));
139    }
140
141    #[test]
142    fn extend_table() {
143        let mut table = Table::new();
144        table.add_row(Row::new(vec![Cell::new("ABC"), Cell::new("DEFG"), Cell::new("HIJKLMN")]));
145        table.extend(vec![vec!["A", "B", "C"]]);
146        let t2 = table.clone();
147        table.extend(t2.rows);
148        assert_eq!(table.get_row(1).unwrap().get_cell(2).unwrap().get_content(), "C");
149        assert_eq!(table.get_row(2).unwrap().get_cell(1).unwrap().get_content(), "DEFG");
150    }
151}