table_extractor/writer/
csv.rs

1use crate::error::Result;
2use crate::{Table, Writer};
3use csv::WriterBuilder;
4use std::io::Write as IoWrite;
5
6pub struct CsvWriter;
7
8impl CsvWriter {
9    pub fn new() -> Self {
10        Self
11    }
12}
13
14impl Default for CsvWriter {
15    fn default() -> Self {
16        Self::new()
17    }
18}
19
20impl Writer for CsvWriter {
21    fn write(&self, table: &Table, output: &mut dyn IoWrite) -> Result<()> {
22        // Write directly to output instead of buffering in Vec
23        // The csv crate uses an internal buffer, and stdout is already wrapped in BufWriter
24        let mut writer = WriterBuilder::new().has_headers(false).from_writer(output);
25
26        // Write headers
27        writer.write_record(&table.headers)?;
28
29        // Write rows
30        for row in &table.rows {
31            writer.write_record(row)?;
32        }
33
34        // Flush the csv writer to ensure all data is written
35        writer.flush()?;
36
37        Ok(())
38    }
39}
40
41#[cfg(test)]
42mod tests {
43    use super::*;
44
45    #[test]
46    fn test_write_csv() {
47        let table = Table::new(
48            vec!["id".to_string(), "name".to_string()],
49            vec![
50                vec!["1".to_string(), "Alice".to_string()],
51                vec!["2".to_string(), "Bob".to_string()],
52            ],
53        );
54
55        let writer = CsvWriter::new();
56        let mut output = Vec::new();
57        writer.write(&table, &mut output).unwrap();
58
59        let result = String::from_utf8(output).unwrap();
60        assert_eq!(result, "id,name\n1,Alice\n2,Bob\n");
61    }
62
63    #[test]
64    fn test_write_csv_with_quotes() {
65        let table = Table::new(
66            vec!["id".to_string(), "name".to_string()],
67            vec![vec![
68                "1".to_string(),
69                "Alice, Bob".to_string(), // Contains comma, should be quoted
70            ]],
71        );
72
73        let writer = CsvWriter::new();
74        let mut output = Vec::new();
75        writer.write(&table, &mut output).unwrap();
76
77        let result = String::from_utf8(output).unwrap();
78        assert_eq!(result, "id,name\n1,\"Alice, Bob\"\n");
79    }
80}
table_extractor/writer/csv.rs

table_extractor/writer/
csv.rs
