1use std::fmt::Write as _;
48use std::fs::File;
49use std::io::{BufWriter, Write};
50use std::path::Path;
51
52use crate::{Error, Result};
53
54#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Hash)]
56pub enum HeaderStyle {
57 #[default]
60 Plain,
61 Commented,
64 None,
66}
67
68#[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 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 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 pub fn comment(mut self, text: impl Into<String>) -> Self {
108 self.comments.push(text.into());
109 self
110 }
111
112 pub fn delimiter(mut self, delimiter: impl Into<String>) -> Self {
114 self.delimiter = delimiter.into();
115 self
116 }
117
118 pub fn comment_prefix(mut self, prefix: impl Into<String>) -> Self {
120 self.comment_prefix = prefix.into();
121 self
122 }
123
124 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 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 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}