Skip to main content

ems/io/
csv.rs

1//! Delimited text writer for time series and frequency series.
2//!
3//! The output is a table of `f64` columns:
4//!
5//! ```text
6//! % probe: port1 voltage
7//! % dt = 1e-12
8//! t,u
9//! 0e0,0e0
10//! 1e-12,2.5e-3
11//! ```
12//!
13//! * optional metadata lines first, each starting with a configurable comment
14//!   prefix (`%` by default, `#` is the other common choice);
15//! * a header row with the column names (plain, comment-prefixed, or
16//!   omitted — see [`HeaderStyle`]);
17//! * one row per sample, values separated by the delimiter (`,` by default).
18//!
19//! Values are written in Rust's shortest round-trip exponent notation
20//! (`{:e}`, e.g. `1.5e-9`), so parsing a value back with
21//! `str::parse::<f64>()` returns exactly the written `f64`. Non-finite values
22//! are written as `NaN`, `inf` and `-inf`.
23//!
24//! Column names are written as given (no quoting); avoid names containing the
25//! delimiter or line breaks.
26//!
27//! # Example
28//!
29//! ```
30//! use ems::io::CsvWriter;
31//!
32//! let t = [0.0, 1e-12, 2e-12];
33//! let u = [0.0, 0.5, -0.25];
34//! let mut out = Vec::new();
35//! CsvWriter::new()
36//!     .comment("port 1 voltage")
37//!     .column("t", &t)
38//!     .column("u", &u)
39//!     .write_to(&mut out)?;
40//! assert_eq!(
41//!     String::from_utf8(out).unwrap(),
42//!     "% port 1 voltage\nt,u\n0e0,0e0\n1e-12,5e-1\n2e-12,-2.5e-1\n"
43//! );
44//! # Ok::<(), ems::Error>(())
45//! ```
46
47use std::fmt::Write as _;
48use std::fs::File;
49use std::io::{BufWriter, Write};
50use std::path::Path;
51
52use crate::{Error, Result};
53
54/// How the header row (column names) is written.
55#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Hash)]
56pub enum HeaderStyle {
57    /// A plain row of names, e.g. `t,u` (the default; what spreadsheet and
58    /// CSV readers expect).
59    #[default]
60    Plain,
61    /// The row of names behind the comment prefix, e.g. `% t,u`, so tools
62    /// that skip comment lines see a purely numeric table.
63    Commented,
64    /// No header row.
65    None,
66}
67
68/// Builder for a delimited table of `f64` columns. See the
69/// [module documentation](self).
70#[derive(Debug, Clone)]
71pub struct CsvWriter<'a> {
72    columns: Vec<(String, &'a [f64])>,
73    comments: Vec<String>,
74    delimiter: String,
75    comment_prefix: String,
76    header: HeaderStyle,
77}
78
79impl Default for CsvWriter<'_> {
80    fn default() -> Self {
81        Self::new()
82    }
83}
84
85impl<'a> CsvWriter<'a> {
86    /// Creates an empty table with delimiter `,`, comment prefix `%` and a
87    /// plain header row.
88    pub fn new() -> Self {
89        CsvWriter {
90            columns: Vec::new(),
91            comments: Vec::new(),
92            delimiter: ",".into(),
93            comment_prefix: "%".into(),
94            header: HeaderStyle::Plain,
95        }
96    }
97
98    /// Appends a column. All columns must have the same length.
99    pub fn column(mut self, name: impl Into<String>, values: &'a [f64]) -> Self {
100        self.columns.push((name.into(), values));
101        self
102    }
103
104    /// Adds a metadata line, written before the header behind the comment
105    /// prefix and a space. A text containing line breaks becomes several
106    /// comment lines.
107    pub fn comment(mut self, text: impl Into<String>) -> Self {
108        self.comments.push(text.into());
109        self
110    }
111
112    /// Sets the value delimiter (default `,`; e.g. `"\t"` or `" "`).
113    pub fn delimiter(mut self, delimiter: impl Into<String>) -> Self {
114        self.delimiter = delimiter.into();
115        self
116    }
117
118    /// Sets the comment prefix (default `%`; e.g. `#`).
119    pub fn comment_prefix(mut self, prefix: impl Into<String>) -> Self {
120        self.comment_prefix = prefix.into();
121        self
122    }
123
124    /// Sets how the header row is written (default [`HeaderStyle::Plain`]).
125    pub fn header(mut self, style: HeaderStyle) -> Self {
126        self.header = style;
127        self
128    }
129
130    fn validate(&self) -> Result<()> {
131        if self.delimiter.is_empty() {
132            return Err(Error::Invalid("CSV delimiter is empty".into()));
133        }
134        if (!self.comments.is_empty() || self.header == HeaderStyle::Commented)
135            && self.comment_prefix.is_empty()
136        {
137            return Err(Error::Invalid("CSV comment prefix is empty".into()));
138        }
139        if let Some((first, rest)) = self.columns.split_first() {
140            let n = first.1.len();
141            if let Some((name, v)) = rest.iter().find(|(_, v)| v.len() != n) {
142                return Err(Error::Invalid(format!(
143                    "CSV column {name:?} has {} rows, column {:?} has {n}",
144                    v.len(),
145                    first.0
146                )));
147            }
148        }
149        Ok(())
150    }
151
152    /// Writes the table to `w`. Fails with [`Error::Invalid`] if the columns
153    /// differ in length or the delimiter (or a needed comment prefix) is
154    /// empty; nothing is written then.
155    pub fn write_to<W: Write>(&self, mut w: W) -> Result<()> {
156        self.validate()?;
157        let mut line = String::new();
158        for text in &self.comments {
159            for l in text.lines() {
160                line.clear();
161                line.push_str(&self.comment_prefix);
162                if !l.is_empty() {
163                    line.push(' ');
164                    line.push_str(l);
165                }
166                line.push('\n');
167                w.write_all(line.as_bytes())?;
168            }
169        }
170        if self.header != HeaderStyle::None && !self.columns.is_empty() {
171            line.clear();
172            if self.header == HeaderStyle::Commented {
173                line.push_str(&self.comment_prefix);
174                line.push(' ');
175            }
176            for (n, (name, _)) in self.columns.iter().enumerate() {
177                if n > 0 {
178                    line.push_str(&self.delimiter);
179                }
180                line.push_str(name);
181            }
182            line.push('\n');
183            w.write_all(line.as_bytes())?;
184        }
185        let rows = self.columns.first().map_or(0, |c| c.1.len());
186        for r in 0..rows {
187            line.clear();
188            for (n, (_, v)) in self.columns.iter().enumerate() {
189                if n > 0 {
190                    line.push_str(&self.delimiter);
191                }
192                write!(line, "{:e}", v[r]).ok();
193            }
194            line.push('\n');
195            w.write_all(line.as_bytes())?;
196        }
197        Ok(())
198    }
199
200    /// Writes the table to `path` (buffered).
201    pub fn write(&self, path: impl AsRef<Path>) -> Result<()> {
202        self.validate()?;
203        let mut w = BufWriter::new(File::create(path)?);
204        self.write_to(&mut w)?;
205        w.flush()?;
206        Ok(())
207    }
208}
209
210#[cfg(test)]
211mod tests {
212    use super::*;
213
214    fn render(w: &CsvWriter<'_>) -> String {
215        let mut out = Vec::new();
216        w.write_to(&mut out).unwrap();
217        String::from_utf8(out).unwrap()
218    }
219
220    #[test]
221    fn basic_layout() {
222        let f = [1e9, 2e9];
223        let re = [0.5, -1.0];
224        let im = [0.0, 3.25e-7];
225        let s = render(
226            &CsvWriter::new()
227                .comment_prefix("#")
228                .comment("S11\nZ0 = 50")
229                .delimiter("\t")
230                .column("f", &f)
231                .column("re", &re)
232                .column("im", &im),
233        );
234        assert_eq!(
235            s,
236            "# S11\n# Z0 = 50\nf\tre\tim\n1e9\t5e-1\t0e0\n2e9\t-1e0\t3.25e-7\n"
237        );
238    }
239
240    #[test]
241    fn header_styles() {
242        let t = [1.0];
243        let base = CsvWriter::new().column("t", &t);
244        assert_eq!(render(&base.clone()), "t\n1e0\n");
245        assert_eq!(
246            render(&base.clone().header(HeaderStyle::Commented)),
247            "% t\n1e0\n"
248        );
249        assert_eq!(render(&base.header(HeaderStyle::None)), "1e0\n");
250    }
251
252    #[test]
253    fn float_round_trip() {
254        let vals = [
255            0.0,
256            -0.0,
257            1.0,
258            0.1,
259            1.0 / 3.0,
260            std::f64::consts::PI,
261            f64::MIN_POSITIVE,
262            f64::MAX,
263            -f64::EPSILON,
264            5e-324,
265            1.234_567_890_123_456_7e-300,
266            299_792_458.0,
267        ];
268        let s = render(
269            &CsvWriter::new()
270                .header(HeaderStyle::None)
271                .column("v", &vals),
272        );
273        let back: Vec<f64> = s.lines().map(|l| l.parse().unwrap()).collect();
274        assert_eq!(back.len(), vals.len());
275        for (a, b) in vals.iter().zip(&back) {
276            assert_eq!(a.to_bits(), b.to_bits(), "{a:e} vs {b:e}");
277        }
278    }
279
280    #[test]
281    fn non_finite() {
282        let v = [f64::NAN, f64::INFINITY, f64::NEG_INFINITY];
283        let s = render(&CsvWriter::new().header(HeaderStyle::None).column("v", &v));
284        assert_eq!(s, "NaN\ninf\n-inf\n");
285        let back: Vec<f64> = s.lines().map(|l| l.parse().unwrap()).collect();
286        assert!(back[0].is_nan() && back[1] == f64::INFINITY && back[2] == f64::NEG_INFINITY);
287    }
288
289    #[test]
290    fn errors() {
291        let a = [1.0, 2.0];
292        let b = [1.0];
293        let mut out = Vec::new();
294        let e = CsvWriter::new()
295            .column("a", &a)
296            .column("b", &b)
297            .write_to(&mut out);
298        assert!(matches!(e, Err(Error::Invalid(_))));
299        assert!(out.is_empty());
300        let e = CsvWriter::new()
301            .delimiter("")
302            .column("a", &a)
303            .write_to(&mut out);
304        assert!(matches!(e, Err(Error::Invalid(_))));
305        let e = CsvWriter::new()
306            .comment_prefix("")
307            .comment("x")
308            .write_to(&mut out);
309        assert!(matches!(e, Err(Error::Invalid(_))));
310    }
311
312    #[test]
313    fn write_to_path() {
314        let dir = std::env::temp_dir().join(format!("ems-csv-test-{}", std::process::id()));
315        std::fs::create_dir_all(&dir).unwrap();
316        let path = dir.join("probe.csv");
317        let t = [0.0, 1.0];
318        CsvWriter::new().column("t", &t).write(&path).unwrap();
319        assert_eq!(std::fs::read_to_string(&path).unwrap(), "t\n0e0\n1e0\n");
320        std::fs::remove_dir_all(&dir).unwrap();
321    }
322}