Skip to main content

xil_core/
pycsv.rs

1//! `csv.writer` / `csv.DictWriter` the way CPython writes them.
2//!
3//! Default dialect: `,` delimiter, `"` quote doubled on escape,
4//! QUOTE_MINIMAL, `\r\n` line terminator. A field is quoted when it
5//! contains the delimiter, the quote character, `\r` or `\n`. A row that is
6//! a single empty field is written as `""` so it is not an empty line.
7//!
8//! Values are rendered with Python `str()` semantics via [`cell`]:
9//! `None` → empty, `True`/`False`, ints plain, floats by `repr`.
10
11use std::io::{self, Write};
12
13use serde_json::Value;
14
15use crate::pyjson::float_repr;
16
17/// Python `str()` of a JSON scalar as csv would stringify it.
18pub fn cell(v: &Value) -> String {
19    match v {
20        Value::Null => String::new(),
21        Value::Bool(true) => "True".to_string(),
22        Value::Bool(false) => "False".to_string(),
23        Value::Number(n) => {
24            if let Some(i) = n.as_i64() {
25                i.to_string()
26            } else if let Some(u) = n.as_u64() {
27                u.to_string()
28            } else {
29                float_repr(n.as_f64().unwrap_or(0.0))
30            }
31        }
32        // A non-finite float arrives carrying pyjson's raw marker. The
33        // csv writer calls str() on it, which spells them -inf / inf /
34        // nan — not the Infinity / NaN tokens the json module uses.
35        Value::String(s) => match s.strip_prefix(crate::pyjson::RAW_MARKER) {
36            Some("Infinity") => "inf".to_string(),
37            Some("-Infinity") => "-inf".to_string(),
38            Some("NaN") => "nan".to_string(),
39            Some(other) => other.to_string(),
40            None => s.clone(),
41        },
42        // Lists/dicts are rare in a CSV cell; Python would print their repr.
43        // Close enough for the audit exports this serves.
44        other => crate::pyjson::dumps(other, crate::pyjson::Style::COMPACT),
45    }
46}
47
48fn needs_quotes(field: &str) -> bool {
49    field.contains(',') || field.contains('"') || field.contains('\r') || field.contains('\n')
50}
51
52/// Write one row. `fields` are already stringified.
53pub fn write_row<W: Write>(out: &mut W, fields: &[String]) -> io::Result<()> {
54    let mut line = String::new();
55    for (i, f) in fields.iter().enumerate() {
56        if i > 0 {
57            line.push(',');
58        }
59        if needs_quotes(f) || (fields.len() == 1 && f.is_empty()) {
60            line.push('"');
61            line.push_str(&f.replace('"', "\"\""));
62            line.push('"');
63        } else {
64            line.push_str(f);
65        }
66    }
67    line.push_str("\r\n");
68    out.write_all(line.as_bytes())
69}
70
71/// `DictWriter(fieldnames=cols)`: header, then each row's values in column
72/// order. A missing key writes an empty cell (`restval=""`); extra keys are
73/// ignored (`extrasaction="ignore"`).
74pub fn write_dicts<W: Write>(
75    out: &mut W,
76    cols: &[&str],
77    rows: &[serde_json::Map<String, Value>],
78) -> io::Result<()> {
79    write_row(out, &cols.iter().map(|c| c.to_string()).collect::<Vec<_>>())?;
80    for row in rows {
81        let fields: Vec<String> = cols
82            .iter()
83            .map(|c| row.get(*c).map(cell).unwrap_or_default())
84            .collect();
85        write_row(out, &fields)?;
86    }
87    Ok(())
88}
89
90/// `csv.reader` on a file opened with `newline=""`: the C state machine
91/// from `_csv.c`, default dialect. Quoted fields may span lines and keep
92/// their line endings verbatim; a quote inside an unquoted field is
93/// literal; a doubled quote inside a quoted field is one quote. A blank
94/// line yields an empty row.
95pub fn read_rows(text: &str) -> Vec<Vec<String>> {
96    #[derive(PartialEq)]
97    enum S {
98        StartRecord,
99        StartField,
100        InField,
101        InQuoted,
102        QuoteInQuoted,
103        EatCrnl,
104    }
105    struct Machine {
106        rows: Vec<Vec<String>>,
107        fields: Vec<String>,
108        field: String,
109        state: S,
110    }
111    impl Machine {
112        fn save(&mut self) {
113            self.fields.push(std::mem::take(&mut self.field));
114        }
115        /// The sentinel the C reader feeds after every physical line. A
116        /// record closes when it leaves the machine in StartRecord.
117        fn eol(&mut self) {
118            match self.state {
119                S::StartField | S::InField | S::QuoteInQuoted => {
120                    self.save();
121                    self.state = S::StartRecord;
122                }
123                S::EatCrnl => self.state = S::StartRecord,
124                S::InQuoted => return, // the field continues on the next line
125                S::StartRecord => {}
126            }
127            self.rows.push(std::mem::take(&mut self.fields));
128        }
129    }
130    let mut m = Machine {
131        rows: Vec::new(),
132        fields: Vec::new(),
133        field: String::new(),
134        state: S::StartRecord,
135    };
136
137    // Lines split the way `newline=""` does — on \n, \r or \r\n, with the
138    // terminator kept as ordinary characters.
139    let chars: Vec<char> = text.chars().collect();
140    let mut i = 0;
141    while i < chars.len() {
142        let c = chars[i];
143        let Machine {
144            ref mut fields,
145            ref mut field,
146            ref mut state,
147            ..
148        } = m;
149        match state {
150            S::StartRecord => {
151                if c == '\n' || c == '\r' {
152                    *state = S::EatCrnl;
153                } else {
154                    *state = S::StartField;
155                    continue; // re-dispatch this char
156                }
157            }
158            S::StartField => {
159                if c == '\n' || c == '\r' {
160                    fields.push(std::mem::take(field));
161                    *state = S::EatCrnl;
162                } else if c == '"' {
163                    *state = S::InQuoted;
164                } else if c == ',' {
165                    fields.push(std::mem::take(field));
166                } else {
167                    field.push(c);
168                    *state = S::InField;
169                }
170            }
171            S::InField => {
172                if c == '\n' || c == '\r' {
173                    fields.push(std::mem::take(field));
174                    *state = S::EatCrnl;
175                } else if c == ',' {
176                    fields.push(std::mem::take(field));
177                    *state = S::StartField;
178                } else {
179                    field.push(c);
180                }
181            }
182            S::InQuoted => {
183                if c == '"' {
184                    *state = S::QuoteInQuoted;
185                } else {
186                    field.push(c);
187                }
188            }
189            S::QuoteInQuoted => {
190                if c == '"' {
191                    field.push('"');
192                    *state = S::InQuoted;
193                } else if c == ',' {
194                    fields.push(std::mem::take(field));
195                    *state = S::StartField;
196                } else if c == '\n' || c == '\r' {
197                    fields.push(std::mem::take(field));
198                    *state = S::EatCrnl;
199                } else {
200                    // Non-strict dialect: text after a closing quote is kept.
201                    field.push(c);
202                    *state = S::InField;
203                }
204            }
205            S::EatCrnl => {
206                // Anything but a line ending here is a Python csv.Error;
207                // a file this crate wrote never produces one.
208            }
209        }
210        i += 1;
211        if c == '\n' || (c == '\r' && chars.get(i) != Some(&'\n')) {
212            m.eol();
213        }
214    }
215    // A final line with no terminator still gets its sentinel.
216    if !text.is_empty() && !text.ends_with('\n') && !text.ends_with('\r') {
217        m.eol();
218    }
219    // End of input inside a quoted field: the non-strict reader keeps
220    // what it has.
221    if m.state == S::InQuoted {
222        m.save();
223        m.rows.push(std::mem::take(&mut m.fields));
224    }
225    m.rows
226}
227
228/// `csv.DictReader`: the first row names the columns. A short row leaves
229/// its missing columns absent (Python's `None`); a long row's extras are
230/// dropped; a blank line is skipped.
231pub fn read_dicts(text: &str) -> Vec<indexmap::IndexMap<String, String>> {
232    let mut rows = read_rows(text).into_iter();
233    let Some(header) = rows.next() else {
234        return Vec::new();
235    };
236    rows.filter(|r| !r.is_empty())
237        .map(|r| header.iter().cloned().zip(r).collect())
238        .collect()
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244    use serde_json::json;
245
246    fn rows(text: &str) -> Vec<Vec<&'static str>> {
247        read_rows(text)
248            .into_iter()
249            .map(|r| {
250                r.into_iter()
251                    .map(|s| &*Box::leak(s.into_boxed_str()))
252                    .collect()
253            })
254            .collect()
255    }
256
257    /// Every case here was run through CPython's csv module; the expected
258    /// values are what it returned.
259    #[test]
260    fn reader_matches_cpython_state_machine() {
261        assert_eq!(rows("a,b\r\n1,2\r\n"), vec![vec!["a", "b"], vec!["1", "2"]]);
262        assert_eq!(
263            rows("a,b\r\n\"x,y\",\"q\"\"r\"\r\n"),
264            vec![vec!["a", "b"], vec!["x,y", "q\"r"]]
265        );
266        assert_eq!(
267            rows("a,b\r\n\"multi\nline\",z\r\n"),
268            vec![vec!["a", "b"], vec!["multi\nline", "z"]]
269        );
270        assert_eq!(
271            rows("a,b\r\n\"a\r\nb\",c\r\n"),
272            vec![vec!["a", "b"], vec!["a\r\nb", "c"]]
273        );
274        assert_eq!(
275            rows("a,b\r\n\r\n1,2\r\n"),
276            vec![vec!["a", "b"], vec![], vec!["1", "2"]]
277        );
278        assert_eq!(rows("a,b\r\n1\r\n"), vec![vec!["a", "b"], vec!["1"]]);
279        assert_eq!(rows("a,b\n1,2\n"), vec![vec!["a", "b"], vec!["1", "2"]]);
280        assert_eq!(rows("a,b\r1,2\r"), vec![vec!["a", "b"], vec!["1", "2"]]);
281        assert_eq!(rows("a,b\r\n1,2"), vec![vec!["a", "b"], vec!["1", "2"]]);
282        assert_eq!(
283            rows("a,b\r\nab\"c,d\r\n"),
284            vec![vec!["a", "b"], vec!["ab\"c", "d"]]
285        );
286        assert_eq!(
287            rows("a,b\r\n\"ab\"c,d\r\n"),
288            vec![vec!["a", "b"], vec!["abc", "d"]]
289        );
290        assert_eq!(rows("a,b\r\n\"\",\r\n"), vec![vec!["a", "b"], vec!["", ""]]);
291        assert_eq!(rows("a,b\r\n ,\r\n"), vec![vec!["a", "b"], vec![" ", ""]]);
292        assert_eq!(
293            rows("a,b\r\n1,\"2\"\r\n\r\n"),
294            vec![vec!["a", "b"], vec!["1", "2"], vec![]]
295        );
296        assert_eq!(rows(""), Vec::<Vec<&str>>::new());
297    }
298
299    #[test]
300    fn dict_reader_short_long_and_blank_rows() {
301        let d = read_dicts("a,b\r\n1\r\n\r\n1,2,3\r\n");
302        assert_eq!(d.len(), 2);
303        assert_eq!(d[0].get("a").map(String::as_str), Some("1"));
304        assert_eq!(d[0].get("b"), None, "short row: column absent, like None");
305        assert_eq!(d[1].get("b").map(String::as_str), Some("2"));
306        assert_eq!(d[1].len(), 2, "extras under None are dropped");
307    }
308
309    #[test]
310    fn writer_output_reads_back_verbatim() {
311        let mut buf = Vec::new();
312        let fields = ["plain", "with,comma", "with \"quote\"", "multi\nline", ""];
313        write_row(&mut buf, &fields.map(String::from)).unwrap();
314        let back = read_rows(std::str::from_utf8(&buf).unwrap());
315        assert_eq!(back, vec![fields.map(String::from).to_vec()]);
316    }
317
318    fn row(fields: &[&str]) -> String {
319        let mut buf = Vec::new();
320        write_row(
321            &mut buf,
322            &fields.iter().map(|s| s.to_string()).collect::<Vec<_>>(),
323        )
324        .unwrap();
325        String::from_utf8(buf).unwrap()
326    }
327
328    #[test]
329    fn quoting_matches_python_minimal() {
330        assert_eq!(
331            row(&["a", "b,c", "d\"e", "f\ng", "", "x y", " lead"]),
332            "a,\"b,c\",\"d\"\"e\",\"f\ng\",,x y, lead\r\n"
333        );
334        assert_eq!(row(&[""]), "\"\"\r\n");
335        assert_eq!(row(&["", ""]), ",\r\n");
336    }
337
338    #[test]
339    fn cells_render_like_python_str() {
340        assert_eq!(cell(&json!(5)), "5");
341        assert_eq!(cell(&json!(5.0)), "5.0");
342        assert_eq!(cell(&json!(0.3)), "0.3");
343        assert_eq!(cell(&json!(1e-5)), "1e-05");
344        assert_eq!(cell(&json!(true)), "True");
345        assert_eq!(cell(&Value::Null), "");
346        assert_eq!(cell(&json!("é")), "é");
347    }
348
349    #[test]
350    fn non_finite_floats_use_python_str_spelling() {
351        use crate::pyjson::py_float;
352        assert_eq!(cell(&py_float(f64::NEG_INFINITY)), "-inf");
353        assert_eq!(cell(&py_float(f64::INFINITY)), "inf");
354        assert_eq!(cell(&py_float(f64::NAN)), "nan");
355        assert_eq!(cell(&py_float(-45.45)), "-45.45");
356    }
357
358    #[test]
359    fn dictwriter_fills_missing_and_ignores_extra() {
360        let mut buf = Vec::new();
361        let r1: serde_json::Map<String, Value> = json!({"a": 1, "b": "x", "zzz": 9})
362            .as_object()
363            .unwrap()
364            .clone();
365        let r2: serde_json::Map<String, Value> = json!({"b": "y"}).as_object().unwrap().clone();
366        write_dicts(&mut buf, &["a", "b"], &[r1, r2]).unwrap();
367        assert_eq!(String::from_utf8(buf).unwrap(), "a,b\r\n1,x\r\n,y\r\n");
368    }
369}