1use std::io::{self, Write};
12
13use serde_json::Value;
14
15use crate::pyjson::float_repr;
16
17pub 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 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 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
52pub 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
71pub 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
90pub 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 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, 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 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; }
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 field.push(c);
202 *state = S::InField;
203 }
204 }
205 S::EatCrnl => {
206 }
209 }
210 i += 1;
211 if c == '\n' || (c == '\r' && chars.get(i) != Some(&'\n')) {
212 m.eol();
213 }
214 }
215 if !text.is_empty() && !text.ends_with('\n') && !text.ends_with('\r') {
217 m.eol();
218 }
219 if m.state == S::InQuoted {
222 m.save();
223 m.rows.push(std::mem::take(&mut m.fields));
224 }
225 m.rows
226}
227
228pub 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 #[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}