ggetrs_string/
utils.rs

1use std::fs::File;
2use std::io::BufWriter;
3use std::io::Write;
4use std::str::FromStr;
5
6use anyhow::Result;
7use polars_core::prelude::*;
8use polars_io::prelude::*;
9
10/// Match the output to either a file or stdout
11pub fn match_output(path: Option<String>) -> Result<Box<dyn Write>> {
12    if let Some(path) = path {
13        let file = File::create(path)?;
14        let buffer = BufWriter::new(file);
15        Ok(Box::new(buffer))
16    } else {
17        let file = std::io::stdout();
18        let buffer = BufWriter::new(file);
19        Ok(Box::new(buffer))
20    }
21}
22
23pub enum WriteConfig {
24    CSV { delimiter: u8, header: bool },
25    JSON { format: JsonFormat },
26}
27
28pub fn write_dataframe<W: Write>(
29    output: W,
30    dataframe: &mut DataFrame,
31    config: WriteConfig,
32) -> Result<(), PolarsError> {
33    match config {
34        WriteConfig::CSV { delimiter, header } => CsvWriter::new(output)
35            .with_separator(delimiter)
36            .include_header(header)
37            .finish(dataframe),
38        WriteConfig::JSON { format } => JsonWriter::new(output)
39            .with_json_format(format)
40            .finish(dataframe),
41    }
42}
43
44#[derive(Clone, Copy, Debug)]
45pub enum OutputFormat {
46    Csv,
47    Tsv,
48    Json,
49    Ndjson,
50}
51impl From<OutputFormat> for WriteConfig {
52    fn from(format: OutputFormat) -> Self {
53        match format {
54            OutputFormat::Csv => WriteConfig::CSV {
55                delimiter: b',',
56                header: true,
57            },
58            OutputFormat::Tsv => WriteConfig::CSV {
59                delimiter: b'\t',
60                header: true,
61            },
62            OutputFormat::Json => WriteConfig::JSON {
63                format: JsonFormat::Json,
64            },
65            OutputFormat::Ndjson => WriteConfig::JSON {
66                format: JsonFormat::JsonLines,
67            },
68        }
69    }
70}
71impl FromStr for OutputFormat {
72    type Err = String;
73
74    fn from_str(s: &str) -> Result<Self, Self::Err> {
75        match s.to_lowercase().as_str() {
76            "csv" => Ok(Self::Csv),
77            "tsv" => Ok(Self::Tsv),
78            "json" => Ok(Self::Json),
79            "ndjson" => Ok(Self::Ndjson),
80            _ => Err(format!("Invalid output format: {}", s)),
81        }
82    }
83}
ggetrs_string/utils.rs

ggetrs_string/
utils.rs
