Skip to main content

robinpath_modules/modules/
collection_mod.rs

1use robinpath::{RobinPath, Value};
2
3pub fn register(rp: &mut RobinPath) {
4    rp.register_builtin("collection.pluck", |args, _| {
5        let arr = args.first().cloned().unwrap_or(Value::Null);
6        let key = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
7        if let Value::Array(items) = &arr {
8            let values: Vec<Value> = items
9                .iter()
10                .map(|item| {
11                    if let Value::Object(obj) = item {
12                        obj.get(&key).cloned().unwrap_or(Value::Null)
13                    } else {
14                        Value::Null
15                    }
16                })
17                .collect();
18            Ok(Value::Array(values))
19        } else {
20            Ok(Value::Array(vec![]))
21        }
22    });
23
24    rp.register_builtin("collection.where", |args, _| {
25        let arr = args.first().cloned().unwrap_or(Value::Null);
26        let key = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
27        let value = args.get(2).cloned().unwrap_or(Value::Null);
28        filter_by_key(&arr, &key, |v| v.deep_eq(&value))
29    });
30
31    rp.register_builtin("collection.whereNot", |args, _| {
32        let arr = args.first().cloned().unwrap_or(Value::Null);
33        let key = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
34        let value = args.get(2).cloned().unwrap_or(Value::Null);
35        filter_by_key(&arr, &key, |v| !v.deep_eq(&value))
36    });
37
38    rp.register_builtin("collection.whereGt", |args, _| {
39        let arr = args.first().cloned().unwrap_or(Value::Null);
40        let key = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
41        let threshold = args.get(2).map(|v| v.to_number()).unwrap_or(0.0);
42        filter_by_key(&arr, &key, |v| v.to_number() > threshold)
43    });
44
45    rp.register_builtin("collection.whereLt", |args, _| {
46        let arr = args.first().cloned().unwrap_or(Value::Null);
47        let key = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
48        let threshold = args.get(2).map(|v| v.to_number()).unwrap_or(0.0);
49        filter_by_key(&arr, &key, |v| v.to_number() < threshold)
50    });
51
52    rp.register_builtin("collection.whereGte", |args, _| {
53        let arr = args.first().cloned().unwrap_or(Value::Null);
54        let key = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
55        let threshold = args.get(2).map(|v| v.to_number()).unwrap_or(0.0);
56        filter_by_key(&arr, &key, |v| v.to_number() >= threshold)
57    });
58
59    rp.register_builtin("collection.whereLte", |args, _| {
60        let arr = args.first().cloned().unwrap_or(Value::Null);
61        let key = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
62        let threshold = args.get(2).map(|v| v.to_number()).unwrap_or(0.0);
63        filter_by_key(&arr, &key, |v| v.to_number() <= threshold)
64    });
65
66    rp.register_builtin("collection.sortBy", |args, _| {
67        let arr = args.first().cloned().unwrap_or(Value::Null);
68        let key = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
69        sort_by_key(&arr, &key, false)
70    });
71
72    rp.register_builtin("collection.sortByDesc", |args, _| {
73        let arr = args.first().cloned().unwrap_or(Value::Null);
74        let key = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
75        sort_by_key(&arr, &key, true)
76    });
77
78    rp.register_builtin("collection.unique", |args, _| {
79        let arr = args.first().cloned().unwrap_or(Value::Null);
80        if let Value::Array(items) = &arr {
81            let mut seen = Vec::new();
82            let mut result = Vec::new();
83            for item in items {
84                let key = item.to_json_string();
85                if !seen.contains(&key) {
86                    seen.push(key);
87                    result.push(item.clone());
88                }
89            }
90            Ok(Value::Array(result))
91        } else {
92            Ok(arr)
93        }
94    });
95
96    rp.register_builtin("collection.flatten", |args, _| {
97        let arr = args.first().cloned().unwrap_or(Value::Null);
98        if let Value::Array(items) = &arr {
99            let mut result = Vec::new();
100            for item in items {
101                if let Value::Array(inner) = item {
102                    result.extend(inner.clone());
103                } else {
104                    result.push(item.clone());
105                }
106            }
107            Ok(Value::Array(result))
108        } else {
109            Ok(arr)
110        }
111    });
112
113    rp.register_builtin("collection.reverse", |args, _| {
114        let arr = args.first().cloned().unwrap_or(Value::Null);
115        if let Value::Array(mut items) = arr {
116            items.reverse();
117            Ok(Value::Array(items))
118        } else {
119            Ok(Value::Null)
120        }
121    });
122
123    rp.register_builtin("collection.chunk", |args, _| {
124        let arr = args.first().cloned().unwrap_or(Value::Null);
125        let size = args.get(1).map(|v| v.to_number() as usize).unwrap_or(1).max(1);
126        if let Value::Array(items) = &arr {
127            let chunks: Vec<Value> = items
128                .chunks(size)
129                .map(|c| Value::Array(c.to_vec()))
130                .collect();
131            Ok(Value::Array(chunks))
132        } else {
133            Ok(Value::Array(vec![]))
134        }
135    });
136
137    rp.register_builtin("collection.first", |args, _| {
138        let arr = args.first().cloned().unwrap_or(Value::Null);
139        if let Value::Array(items) = &arr {
140            Ok(items.first().cloned().unwrap_or(Value::Null))
141        } else {
142            Ok(Value::Null)
143        }
144    });
145
146    rp.register_builtin("collection.last", |args, _| {
147        let arr = args.first().cloned().unwrap_or(Value::Null);
148        if let Value::Array(items) = &arr {
149            Ok(items.last().cloned().unwrap_or(Value::Null))
150        } else {
151            Ok(Value::Null)
152        }
153    });
154
155    rp.register_builtin("collection.count", |args, _| {
156        let arr = args.first().cloned().unwrap_or(Value::Null);
157        if let Value::Array(items) = &arr {
158            Ok(Value::Number(items.len() as f64))
159        } else {
160            Ok(Value::Number(0.0))
161        }
162    });
163
164    rp.register_builtin("collection.sum", |args, _| {
165        let arr = args.first().cloned().unwrap_or(Value::Null);
166        let key = args.get(1).map(|v| v.to_display_string());
167        if let Value::Array(items) = &arr {
168            let sum: f64 = items.iter().map(|item| extract_num(item, key.as_deref())).sum();
169            Ok(Value::Number(sum))
170        } else {
171            Ok(Value::Number(0.0))
172        }
173    });
174
175    rp.register_builtin("collection.avg", |args, _| {
176        let arr = args.first().cloned().unwrap_or(Value::Null);
177        let key = args.get(1).map(|v| v.to_display_string());
178        if let Value::Array(items) = &arr {
179            if items.is_empty() { return Ok(Value::Number(0.0)); }
180            let sum: f64 = items.iter().map(|item| extract_num(item, key.as_deref())).sum();
181            Ok(Value::Number(sum / items.len() as f64))
182        } else {
183            Ok(Value::Number(0.0))
184        }
185    });
186
187    rp.register_builtin("collection.min", |args, _| {
188        let arr = args.first().cloned().unwrap_or(Value::Null);
189        let key = args.get(1).map(|v| v.to_display_string());
190        if let Value::Array(items) = &arr {
191            let min = items.iter().map(|item| extract_num(item, key.as_deref())).fold(f64::INFINITY, f64::min);
192            Ok(Value::Number(min))
193        } else {
194            Ok(Value::Null)
195        }
196    });
197
198    rp.register_builtin("collection.max", |args, _| {
199        let arr = args.first().cloned().unwrap_or(Value::Null);
200        let key = args.get(1).map(|v| v.to_display_string());
201        if let Value::Array(items) = &arr {
202            let max = items.iter().map(|item| extract_num(item, key.as_deref())).fold(f64::NEG_INFINITY, f64::max);
203            Ok(Value::Number(max))
204        } else {
205            Ok(Value::Null)
206        }
207    });
208
209    rp.register_builtin("collection.groupBy", |args, _| {
210        let arr = args.first().cloned().unwrap_or(Value::Null);
211        let key = args.get(1).map(|v| v.to_display_string()).unwrap_or_default();
212        if let Value::Array(items) = &arr {
213            let mut groups = indexmap::IndexMap::new();
214            for item in items {
215                let group_key = if let Value::Object(obj) = item {
216                    obj.get(&key).map(|v| v.to_display_string()).unwrap_or_else(|| "null".to_string())
217                } else {
218                    "null".to_string()
219                };
220                groups.entry(group_key).or_insert_with(Vec::new).push(item.clone());
221            }
222            let result: indexmap::IndexMap<String, Value> = groups
223                .into_iter()
224                .map(|(k, v)| (k, Value::Array(v)))
225                .collect();
226            Ok(Value::Object(result))
227        } else {
228            Ok(Value::Object(indexmap::IndexMap::new()))
229        }
230    });
231
232    rp.register_builtin("collection.compact", |args, _| {
233        let arr = args.first().cloned().unwrap_or(Value::Null);
234        if let Value::Array(items) = &arr {
235            let result: Vec<Value> = items.iter().filter(|v| v.is_truthy()).cloned().collect();
236            Ok(Value::Array(result))
237        } else {
238            Ok(arr)
239        }
240    });
241
242    rp.register_builtin("collection.zip", |args, _| {
243        let a = args.first().cloned().unwrap_or(Value::Null);
244        let b = args.get(1).cloned().unwrap_or(Value::Null);
245        if let (Value::Array(arr_a), Value::Array(arr_b)) = (&a, &b) {
246            let result: Vec<Value> = arr_a.iter().zip(arr_b.iter())
247                .map(|(x, y)| Value::Array(vec![x.clone(), y.clone()]))
248                .collect();
249            Ok(Value::Array(result))
250        } else {
251            Ok(Value::Array(vec![]))
252        }
253    });
254
255    rp.register_builtin("collection.difference", |args, _| {
256        let a = args.first().cloned().unwrap_or(Value::Null);
257        let b = args.get(1).cloned().unwrap_or(Value::Null);
258        if let (Value::Array(arr_a), Value::Array(arr_b)) = (&a, &b) {
259            let b_keys: Vec<String> = arr_b.iter().map(|v| v.to_json_string()).collect();
260            let result: Vec<Value> = arr_a.iter().filter(|v| !b_keys.contains(&v.to_json_string())).cloned().collect();
261            Ok(Value::Array(result))
262        } else {
263            Ok(Value::Array(vec![]))
264        }
265    });
266
267    rp.register_builtin("collection.intersection", |args, _| {
268        let a = args.first().cloned().unwrap_or(Value::Null);
269        let b = args.get(1).cloned().unwrap_or(Value::Null);
270        if let (Value::Array(arr_a), Value::Array(arr_b)) = (&a, &b) {
271            let b_keys: Vec<String> = arr_b.iter().map(|v| v.to_json_string()).collect();
272            let result: Vec<Value> = arr_a.iter().filter(|v| b_keys.contains(&v.to_json_string())).cloned().collect();
273            Ok(Value::Array(result))
274        } else {
275            Ok(Value::Array(vec![]))
276        }
277    });
278
279    rp.register_builtin("collection.union", |args, _| {
280        let a = args.first().cloned().unwrap_or(Value::Null);
281        let b = args.get(1).cloned().unwrap_or(Value::Null);
282        if let (Value::Array(arr_a), Value::Array(arr_b)) = (&a, &b) {
283            let mut seen = Vec::new();
284            let mut result = Vec::new();
285            for item in arr_a.iter().chain(arr_b.iter()) {
286                let key = item.to_json_string();
287                if !seen.contains(&key) {
288                    seen.push(key);
289                    result.push(item.clone());
290                }
291            }
292            Ok(Value::Array(result))
293        } else {
294            Ok(Value::Array(vec![]))
295        }
296    });
297
298    rp.register_builtin("collection.take", |args, _| {
299        let arr = args.first().cloned().unwrap_or(Value::Null);
300        let n = args.get(1).map(|v| v.to_number() as usize).unwrap_or(1);
301        if let Value::Array(items) = &arr {
302            Ok(Value::Array(items.iter().take(n).cloned().collect()))
303        } else {
304            Ok(Value::Array(vec![]))
305        }
306    });
307
308    rp.register_builtin("collection.skip", |args, _| {
309        let arr = args.first().cloned().unwrap_or(Value::Null);
310        let n = args.get(1).map(|v| v.to_number() as usize).unwrap_or(0);
311        if let Value::Array(items) = &arr {
312            Ok(Value::Array(items.iter().skip(n).cloned().collect()))
313        } else {
314            Ok(Value::Array(vec![]))
315        }
316    });
317
318    rp.register_builtin("collection.contains", |args, _| {
319        let arr = args.first().cloned().unwrap_or(Value::Null);
320        let value = args.get(1).cloned().unwrap_or(Value::Null);
321        if let Value::Array(items) = &arr {
322            Ok(Value::Bool(items.iter().any(|v| v.deep_eq(&value))))
323        } else {
324            Ok(Value::Bool(false))
325        }
326    });
327
328    rp.register_builtin("collection.indexOf", |args, _| {
329        let arr = args.first().cloned().unwrap_or(Value::Null);
330        let value = args.get(1).cloned().unwrap_or(Value::Null);
331        if let Value::Array(items) = &arr {
332            match items.iter().position(|v| v.deep_eq(&value)) {
333                Some(idx) => Ok(Value::Number(idx as f64)),
334                None => Ok(Value::Number(-1.0)),
335            }
336        } else {
337            Ok(Value::Number(-1.0))
338        }
339    });
340}
341
342fn filter_by_key(arr: &Value, key: &str, pred: impl Fn(&Value) -> bool) -> Result<Value, String> {
343    if let Value::Array(items) = arr {
344        let result: Vec<Value> = items
345            .iter()
346            .filter(|item| {
347                if let Value::Object(obj) = item {
348                    obj.get(key).map_or(false, &pred)
349                } else {
350                    false
351                }
352            })
353            .cloned()
354            .collect();
355        Ok(Value::Array(result))
356    } else {
357        Ok(Value::Array(vec![]))
358    }
359}
360
361fn sort_by_key(arr: &Value, key: &str, desc: bool) -> Result<Value, String> {
362    if let Value::Array(items) = arr {
363        let mut sorted = items.clone();
364        sorted.sort_by(|a, b| {
365            let va = if let Value::Object(obj) = a { obj.get(key).cloned().unwrap_or(Value::Null) } else { Value::Null };
366            let vb = if let Value::Object(obj) = b { obj.get(key).cloned().unwrap_or(Value::Null) } else { Value::Null };
367            let cmp = va.to_number().partial_cmp(&vb.to_number()).unwrap_or(std::cmp::Ordering::Equal);
368            if desc { cmp.reverse() } else { cmp }
369        });
370        Ok(Value::Array(sorted))
371    } else {
372        Ok(Value::Array(vec![]))
373    }
374}
375
376fn extract_num(item: &Value, key: Option<&str>) -> f64 {
377    match key {
378        Some(k) => {
379            if let Value::Object(obj) = item {
380                obj.get(k).map(|v| v.to_number()).unwrap_or(0.0)
381            } else {
382                0.0
383            }
384        }
385        None => item.to_number(),
386    }
387}