jpx_core/extensions/
utility.rs

1//! Utility functions.
2
3use std::collections::HashSet;
4
5use serde_json::{Number, Value};
6
7use crate::functions::Function;
8use crate::interpreter::SearchResult;
9use crate::registry::register_if_enabled;
10use crate::{Context, Runtime, arg, defn};
11
12// =============================================================================
13// now(fallback?) -> number (Unix timestamp in seconds)
14// =============================================================================
15
16defn!(NowFn, vec![], Some(arg!(number)));
17
18impl Function for NowFn {
19    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
20        self.signature.validate(args, ctx)?;
21
22        if let Some(fallback) = args.first()
23            && let Some(n) = fallback.as_f64()
24        {
25            return Ok(Value::Number(
26                Number::from_f64(n).unwrap_or_else(|| Number::from(0)),
27            ));
28        }
29
30        let timestamp = std::time::SystemTime::now()
31            .duration_since(std::time::UNIX_EPOCH)
32            .map(|d| d.as_secs())
33            .unwrap_or(0);
34
35        Ok(Value::Number(Number::from(timestamp)))
36    }
37}
38
39// =============================================================================
40// now_ms(fallback?) -> number (Unix timestamp in milliseconds)
41// =============================================================================
42
43defn!(NowMsFn, vec![], Some(arg!(number)));
44
45impl Function for NowMsFn {
46    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
47        self.signature.validate(args, ctx)?;
48
49        if let Some(fallback) = args.first()
50            && let Some(n) = fallback.as_f64()
51        {
52            return Ok(Value::Number(
53                Number::from_f64(n).unwrap_or_else(|| Number::from(0)),
54            ));
55        }
56
57        let timestamp = std::time::SystemTime::now()
58            .duration_since(std::time::UNIX_EPOCH)
59            .map(|d| d.as_millis() as u64)
60            .unwrap_or(0);
61
62        Ok(Value::Number(Number::from(timestamp)))
63    }
64}
65
66// =============================================================================
67// default(value, default_value) -> value if not null, else default
68// =============================================================================
69
70defn!(DefaultFn, vec![arg!(any), arg!(any)], None);
71
72impl Function for DefaultFn {
73    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
74        self.signature.validate(args, ctx)?;
75
76        if args[0].is_null() {
77            Ok(args[1].clone())
78        } else {
79            Ok(args[0].clone())
80        }
81    }
82}
83
84// =============================================================================
85// if(condition, then_value, else_value) -> any
86// =============================================================================
87
88defn!(IfFn, vec![arg!(any), arg!(any), arg!(any)], None);
89
90impl Function for IfFn {
91    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
92        self.signature.validate(args, ctx)?;
93
94        let condition = &args[0];
95        let then_value = &args[1];
96        let else_value = &args[2];
97
98        let is_truthy = match condition {
99            Value::Bool(b) => *b,
100            Value::Null => false,
101            _ => true,
102        };
103
104        if is_truthy {
105            Ok(then_value.clone())
106        } else {
107            Ok(else_value.clone())
108        }
109    }
110}
111
112// =============================================================================
113// coalesce(...) -> any (first non-null value)
114// =============================================================================
115
116defn!(CoalesceFn, vec![], Some(arg!(any)));
117
118impl Function for CoalesceFn {
119    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
120        self.signature.validate(args, ctx)?;
121        for arg in args {
122            if !arg.is_null() {
123                return Ok(arg.clone());
124            }
125        }
126        Ok(Value::Null)
127    }
128}
129
130// =============================================================================
131// json_encode(any) -> string
132// =============================================================================
133
134defn!(JsonEncodeFn, vec![arg!(any)], None);
135
136impl Function for JsonEncodeFn {
137    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
138        self.signature.validate(args, ctx)?;
139
140        let json_str = serde_json::to_string(&args[0])
141            .map_err(|_| crate::functions::custom_error(ctx, "Failed to encode as JSON"))?;
142
143        Ok(Value::String(json_str))
144    }
145}
146
147// =============================================================================
148// pretty(any, indent?) -> string
149// =============================================================================
150
151defn!(PrettyFn, vec![arg!(any)], Some(arg!(number)));
152
153impl Function for PrettyFn {
154    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
155        self.signature.validate(args, ctx)?;
156
157        let indent = if args.len() > 1 {
158            args[1].as_f64().unwrap_or(2.0) as usize
159        } else {
160            2
161        };
162
163        // For default indent of 2, use built-in to_string_pretty
164        if indent == 2 {
165            let pretty_str = serde_json::to_string_pretty(&args[0])
166                .map_err(|_| crate::functions::custom_error(ctx, "Failed to serialize as JSON"))?;
167            return Ok(Value::String(pretty_str));
168        }
169
170        // For custom indent, manually format
171        let json_str = serde_json::to_string(&args[0])
172            .map_err(|_| crate::functions::custom_error(ctx, "Failed to serialize as JSON"))?;
173
174        let pretty_str = pretty_print_json(&json_str, indent);
175        Ok(Value::String(pretty_str))
176    }
177}
178
179/// Pretty print JSON with custom indentation.
180fn pretty_print_json(json: &str, indent_size: usize) -> String {
181    let mut result = String::new();
182    let mut depth = 0;
183    let mut in_string = false;
184    let mut escape_next = false;
185    let indent = " ".repeat(indent_size);
186
187    for ch in json.chars() {
188        if escape_next {
189            result.push(ch);
190            escape_next = false;
191            continue;
192        }
193
194        if ch == '\\' && in_string {
195            result.push(ch);
196            escape_next = true;
197            continue;
198        }
199
200        if ch == '"' {
201            in_string = !in_string;
202            result.push(ch);
203            continue;
204        }
205
206        if in_string {
207            result.push(ch);
208            continue;
209        }
210
211        match ch {
212            '{' | '[' => {
213                result.push(ch);
214                depth += 1;
215                result.push('\n');
216                for _ in 0..depth {
217                    result.push_str(&indent);
218                }
219            }
220            '}' | ']' => {
221                depth -= 1;
222                result.push('\n');
223                for _ in 0..depth {
224                    result.push_str(&indent);
225                }
226                result.push(ch);
227            }
228            ',' => {
229                result.push(ch);
230                result.push('\n');
231                for _ in 0..depth {
232                    result.push_str(&indent);
233                }
234            }
235            ':' => {
236                result.push_str(": ");
237            }
238            ' ' | '\n' | '\t' | '\r' => {
239                // Skip whitespace in compact JSON
240            }
241            _ => {
242                result.push(ch);
243            }
244        }
245    }
246
247    result
248}
249
250// =============================================================================
251// json_decode(string) -> any
252// =============================================================================
253
254defn!(JsonDecodeFn, vec![arg!(string)], None);
255
256impl Function for JsonDecodeFn {
257    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
258        self.signature.validate(args, ctx)?;
259
260        let s = args[0]
261            .as_str()
262            .ok_or_else(|| crate::functions::custom_error(ctx, "Expected string argument"))?;
263
264        // Return null for invalid JSON instead of erroring
265        match serde_json::from_str::<Value>(s) {
266            Ok(val) => Ok(val),
267            Err(_) => Ok(Value::Null),
268        }
269    }
270}
271
272// =============================================================================
273// json_pointer(any, string) -> any (RFC 6901 JSON Pointer)
274// =============================================================================
275
276defn!(JsonPointerFn, vec![arg!(any), arg!(string)], None);
277
278impl Function for JsonPointerFn {
279    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
280        self.signature.validate(args, ctx)?;
281
282        let pointer = args[1].as_str().ok_or_else(|| {
283            crate::functions::custom_error(ctx, "Expected string pointer argument")
284        })?;
285
286        // Use serde_json's built-in pointer method directly on the Value
287        match args[0].pointer(pointer) {
288            Some(result) => Ok(result.clone()),
289            None => Ok(Value::Null),
290        }
291    }
292}
293
294// =============================================================================
295// env() -> object (all environment variables)
296// =============================================================================
297
298defn!(EnvFn, vec![], None);
299
300impl Function for EnvFn {
301    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
302        self.signature.validate(args, ctx)?;
303
304        let mut map = serde_json::Map::new();
305        for (key, value) in std::env::vars() {
306            map.insert(key, Value::String(value));
307        }
308
309        Ok(Value::Object(map))
310    }
311}
312
313// =============================================================================
314// get_env(name) -> string | null (single environment variable)
315// =============================================================================
316
317defn!(GetEnvFn, vec![arg!(string)], None);
318
319impl Function for GetEnvFn {
320    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
321        self.signature.validate(args, ctx)?;
322
323        let name = args[0]
324            .as_str()
325            .ok_or_else(|| crate::functions::custom_error(ctx, "Expected string argument"))?;
326
327        match std::env::var(name) {
328            Ok(value) => Ok(Value::String(value)),
329            Err(_) => Ok(Value::Null),
330        }
331    }
332}
333
334/// Register utility functions that are in the enabled set.
335pub fn register_filtered(runtime: &mut Runtime, enabled: &HashSet<&str>) {
336    register_if_enabled(runtime, "now", enabled, Box::new(NowFn::new()));
337    register_if_enabled(runtime, "now_ms", enabled, Box::new(NowMsFn::new()));
338    register_if_enabled(runtime, "default", enabled, Box::new(DefaultFn::new()));
339    register_if_enabled(runtime, "if", enabled, Box::new(IfFn::new()));
340    register_if_enabled(runtime, "coalesce", enabled, Box::new(CoalesceFn::new()));
341    register_if_enabled(
342        runtime,
343        "json_encode",
344        enabled,
345        Box::new(JsonEncodeFn::new()),
346    );
347    register_if_enabled(runtime, "to_json", enabled, Box::new(JsonEncodeFn::new()));
348    register_if_enabled(
349        runtime,
350        "json_decode",
351        enabled,
352        Box::new(JsonDecodeFn::new()),
353    );
354    register_if_enabled(
355        runtime,
356        "json_pointer",
357        enabled,
358        Box::new(JsonPointerFn::new()),
359    );
360    register_if_enabled(runtime, "pretty", enabled, Box::new(PrettyFn::new()));
361    register_if_enabled(runtime, "env", enabled, Box::new(EnvFn::new()));
362    register_if_enabled(runtime, "get_env", enabled, Box::new(GetEnvFn::new()));
363}
364
365#[cfg(test)]
366mod tests {
367    use crate::Runtime;
368    use serde_json::json;
369
370    fn setup_runtime() -> Runtime {
371        Runtime::builder()
372            .with_standard()
373            .with_all_extensions()
374            .build()
375    }
376
377    #[test]
378    fn test_default() {
379        let runtime = setup_runtime();
380        let expr = runtime.compile("default(@, 'fallback')").unwrap();
381
382        let result = expr.search(&json!(null)).unwrap();
383        assert_eq!(result.as_str().unwrap(), "fallback");
384
385        let result = expr.search(&json!("value")).unwrap();
386        assert_eq!(result.as_str().unwrap(), "value");
387    }
388
389    #[test]
390    fn test_if() {
391        let runtime = setup_runtime();
392        let expr = runtime.compile("if(`true`, 'yes', 'no')").unwrap();
393        let result = expr.search(&json!(null)).unwrap();
394        assert_eq!(result.as_str().unwrap(), "yes");
395
396        let expr = runtime.compile("if(`false`, 'yes', 'no')").unwrap();
397        let result = expr.search(&json!(null)).unwrap();
398        assert_eq!(result.as_str().unwrap(), "no");
399    }
400
401    #[test]
402    fn test_json_decode_object() {
403        let runtime = setup_runtime();
404        // Test parsing a JSON object string
405        let expr = runtime.compile("json_decode(@)").unwrap();
406        let data = json!(r#"{"a":1,"b":2}"#);
407        let result = expr.search(&data).unwrap();
408        assert!(result.is_object());
409        let obj = result.as_object().unwrap();
410        assert!(obj.contains_key("a"));
411    }
412
413    #[test]
414    fn test_json_decode_from_field() {
415        let runtime = setup_runtime();
416        let expr = runtime.compile("json_decode(s)").unwrap();
417        let data = json!({"s": r#"{"a":1,"b":2}"#});
418
419        let result = expr.search(&data);
420        assert!(result.is_ok());
421        let val = result.unwrap();
422        assert!(val.is_object());
423    }
424
425    #[test]
426    fn test_json_decode_invalid_returns_null() {
427        let runtime = setup_runtime();
428        let expr = runtime.compile("json_decode(@)").unwrap();
429        let data = json!("not valid json");
430        let result = expr.search(&data).unwrap();
431        assert!(result.is_null());
432    }
433
434    #[test]
435    fn test_json_encode() {
436        let runtime = setup_runtime();
437        let expr = runtime.compile("json_encode(@)").unwrap();
438        let data = json!({"a": 1});
439        let result = expr.search(&data).unwrap();
440        assert_eq!(result.as_str().unwrap(), r#"{"a":1}"#);
441    }
442
443    #[test]
444    fn test_json_pointer_nested() {
445        let runtime = setup_runtime();
446        let expr = runtime.compile("json_pointer(@, '/foo/bar/1')").unwrap();
447        let data = json!({"foo": {"bar": [1, 2, 3]}});
448        let result = expr.search(&data).unwrap();
449        assert_eq!(result.as_f64().unwrap(), 2.0);
450    }
451
452    #[test]
453    fn test_json_pointer_root() {
454        let runtime = setup_runtime();
455        let expr = runtime.compile("json_pointer(@, '')").unwrap();
456        let data = json!({"a": 1});
457        let result = expr.search(&data).unwrap();
458        assert!(result.is_object());
459    }
460
461    #[test]
462    fn test_json_pointer_missing() {
463        let runtime = setup_runtime();
464        let expr = runtime.compile("json_pointer(@, '/missing')").unwrap();
465        let data = json!({"a": 1});
466        let result = expr.search(&data).unwrap();
467        assert!(result.is_null());
468    }
469
470    #[test]
471    fn test_json_pointer_array() {
472        let runtime = setup_runtime();
473        let expr = runtime.compile("json_pointer(@, '/0')").unwrap();
474        let data = json!([1, 2, 3]);
475        let result = expr.search(&data).unwrap();
476        assert_eq!(result.as_f64().unwrap(), 1.0);
477    }
478
479    #[test]
480    fn test_pretty_default_indent() {
481        let runtime = setup_runtime();
482        let expr = runtime.compile("pretty(@)").unwrap();
483        let data = json!({"a": 1, "b": [2, 3]});
484        let result = expr.search(&data).unwrap();
485        let s = result.as_str().unwrap();
486        assert!(s.contains('\n'));
487        assert!(s.contains("  ")); // 2-space indent
488    }
489
490    #[test]
491    fn test_pretty_custom_indent() {
492        let runtime = setup_runtime();
493        let expr = runtime.compile("pretty(@, `4`)").unwrap();
494        let data = json!({"a": 1});
495        let result = expr.search(&data).unwrap();
496        let s = result.as_str().unwrap();
497        assert!(s.contains("    ")); // 4-space indent
498    }
499
500    #[test]
501    fn test_pretty_simple_value() {
502        let runtime = setup_runtime();
503        let expr = runtime.compile("pretty(@)").unwrap();
504        let data = json!("hello");
505        let result = expr.search(&data).unwrap();
506        assert_eq!(result.as_str().unwrap(), "\"hello\"");
507    }
508
509    #[test]
510    fn test_env_returns_object() {
511        let runtime = setup_runtime();
512        let expr = runtime.compile("env()").unwrap();
513        let result = expr.search(&json!(null)).unwrap();
514        assert!(result.is_object());
515    }
516
517    #[test]
518    fn test_get_env_existing() {
519        // PATH should exist on all systems
520        let runtime = setup_runtime();
521        let expr = runtime.compile("get_env('PATH')").unwrap();
522        let result = expr.search(&json!(null)).unwrap();
523        assert!(result.is_string());
524    }
525
526    #[test]
527    fn test_get_env_missing() {
528        let runtime = setup_runtime();
529        let expr = runtime
530            .compile("get_env('THIS_ENV_VAR_SHOULD_NOT_EXIST_12345')")
531            .unwrap();
532        let result = expr.search(&json!(null)).unwrap();
533        assert!(result.is_null());
534    }
535}
jpx_core/extensions/utility.rs

jpx_core/extensions/
utility.rs
