prettytable/
csv.rs

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