Skip to main content

robinpath_modules/modules/
csv_mod.rs

1use robinpath::{RobinPath, Value};
2
3pub fn register(rp: &mut RobinPath) {
4    rp.register_builtin("csv.parse", |args, _| {
5        let input = args.first().map(|v| v.to_display_string()).unwrap_or_default();
6        let delimiter = args
7            .get(1)
8            .and_then(|v| v.as_str().and_then(|s| s.chars().next()))
9            .unwrap_or(',');
10        let rows = parse_csv_rows(&input, delimiter);
11        if rows.is_empty() {
12            return Ok(Value::Array(vec![]));
13        }
14        let headers = &rows[0];
15        let result: Vec<Value> = rows[1..]
16            .iter()
17            .map(|row| {
18                let mut obj = indexmap::IndexMap::new();
19                for (i, header) in headers.iter().enumerate() {
20                    let val = row.get(i).cloned().unwrap_or_default();
21                    obj.insert(header.clone(), Value::String(val));
22                }
23                Value::Object(obj)
24            })
25            .collect();
26        Ok(Value::Array(result))
27    });
28
29    rp.register_builtin("csv.stringify", |args, _| {
30        let data = args.first().cloned().unwrap_or(Value::Null);
31        let delimiter = args
32            .get(1)
33            .and_then(|v| v.as_str().and_then(|s| s.chars().next()))
34            .unwrap_or(',');
35        match &data {
36            Value::Array(arr) => {
37                if arr.is_empty() {
38                    return Ok(Value::String(String::new()));
39                }
40                // Collect all headers from first object
41                let headers: Vec<String> = if let Some(Value::Object(first)) = arr.first() {
42                    first.keys().cloned().collect()
43                } else {
44                    return Ok(Value::String(String::new()));
45                };
46                let mut lines = Vec::new();
47                lines.push(
48                    headers
49                        .iter()
50                        .map(|h| csv_escape(h, delimiter))
51                        .collect::<Vec<_>>()
52                        .join(&delimiter.to_string()),
53                );
54                for item in arr {
55                    if let Value::Object(obj) = item {
56                        let row: Vec<String> = headers
57                            .iter()
58                            .map(|h| {
59                                let val = obj
60                                    .get(h)
61                                    .map(|v| v.to_display_string())
62                                    .unwrap_or_default();
63                                csv_escape(&val, delimiter)
64                            })
65                            .collect();
66                        lines.push(row.join(&delimiter.to_string()));
67                    }
68                }
69                Ok(Value::String(lines.join("\n")))
70            }
71            _ => Ok(Value::String(String::new())),
72        }
73    });
74
75    rp.register_builtin("csv.headers", |args, _| {
76        let input = args.first().map(|v| v.to_display_string()).unwrap_or_default();
77        let delimiter = args
78            .get(1)
79            .and_then(|v| v.as_str().and_then(|s| s.chars().next()))
80            .unwrap_or(',');
81        let rows = parse_csv_rows(&input, delimiter);
82        if rows.is_empty() {
83            return Ok(Value::Array(vec![]));
84        }
85        let headers: Vec<Value> = rows[0].iter().map(|h| Value::String(h.clone())).collect();
86        Ok(Value::Array(headers))
87    });
88
89    rp.register_builtin("csv.column", |args, _| {
90        let input = args.first().map(|v| v.to_display_string()).unwrap_or_default();
91        let col_name = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
92        let delimiter = args
93            .get(2)
94            .and_then(|v| v.as_str().and_then(|s| s.chars().next()))
95            .unwrap_or(',');
96        let rows = parse_csv_rows(&input, delimiter);
97        if rows.is_empty() {
98            return Ok(Value::Array(vec![]));
99        }
100        let headers = &rows[0];
101        let col_idx = headers.iter().position(|h| h == &col_name);
102        match col_idx {
103            Some(idx) => {
104                let values: Vec<Value> = rows[1..]
105                    .iter()
106                    .map(|row| {
107                        Value::String(row.get(idx).cloned().unwrap_or_default())
108                    })
109                    .collect();
110                Ok(Value::Array(values))
111            }
112            None => Ok(Value::Array(vec![])),
113        }
114    });
115
116    rp.register_builtin("csv.rows", |args, _| {
117        let input = args.first().map(|v| v.to_display_string()).unwrap_or_default();
118        let delimiter = args
119            .get(1)
120            .and_then(|v| v.as_str().and_then(|s| s.chars().next()))
121            .unwrap_or(',');
122        let rows = parse_csv_rows(&input, delimiter);
123        let result: Vec<Value> = rows
124            .into_iter()
125            .map(|row| Value::Array(row.into_iter().map(Value::String).collect()))
126            .collect();
127        Ok(Value::Array(result))
128    });
129}
130
131fn parse_csv_rows(input: &str, delimiter: char) -> Vec<Vec<String>> {
132    let mut rows = Vec::new();
133    let mut current_row: Vec<String> = Vec::new();
134    let mut current_field = String::new();
135    let mut in_quotes = false;
136    let mut chars = input.chars().peekable();
137
138    while let Some(ch) = chars.next() {
139        if in_quotes {
140            if ch == '"' {
141                if chars.peek() == Some(&'"') {
142                    // Escaped quote
143                    current_field.push('"');
144                    chars.next();
145                } else {
146                    in_quotes = false;
147                }
148            } else {
149                current_field.push(ch);
150            }
151        } else if ch == '"' {
152            in_quotes = true;
153        } else if ch == delimiter {
154            current_row.push(current_field.clone());
155            current_field.clear();
156        } else if ch == '\n' {
157            current_row.push(current_field.clone());
158            current_field.clear();
159            if !current_row.iter().all(|f| f.is_empty()) || current_row.len() > 1 {
160                rows.push(current_row.clone());
161            }
162            current_row.clear();
163        } else if ch == '\r' {
164            // Skip \r, handled with \n
165        } else {
166            current_field.push(ch);
167        }
168    }
169
170    // Last field/row
171    if !current_field.is_empty() || !current_row.is_empty() {
172        current_row.push(current_field);
173        rows.push(current_row);
174    }
175
176    rows
177}
178
179fn csv_escape(s: &str, delimiter: char) -> String {
180    if s.contains(delimiter) || s.contains('"') || s.contains('\n') {
181        format!("\"{}\"", s.replace('"', "\"\""))
182    } else {
183        s.to_string()
184    }
185}