vrl/stdlib/
map_keys.rs

1use crate::compiler::prelude::*;
2
3fn map_keys<T>(
4    value: Value,
5    recursive: bool,
6    ctx: &mut Context,
7    runner: &closure::Runner<T>,
8) -> Resolved
9where
10    T: Fn(&mut Context) -> Resolved,
11{
12    let mut iter = value.into_iter(recursive);
13
14    for item in iter.by_ref() {
15        if let IterItem::KeyValue(key, _) = item {
16            runner.map_key(ctx, key)?;
17        }
18    }
19
20    Ok(iter.into())
21}
22
23#[derive(Clone, Copy, Debug)]
24pub struct MapKeys;
25
26impl Function for MapKeys {
27    fn identifier(&self) -> &'static str {
28        "map_keys"
29    }
30
31    fn parameters(&self) -> &'static [Parameter] {
32        &[
33            Parameter {
34                keyword: "value",
35                kind: kind::OBJECT,
36                required: true,
37            },
38            Parameter {
39                keyword: "recursive",
40                kind: kind::BOOLEAN,
41                required: false,
42            },
43        ]
44    }
45
46    fn examples(&self) -> &'static [Example] {
47        &[
48            Example {
49                title: "map object keys",
50                source: r#"map_keys({ "a": 1, "b": 2 }) -> |key| { upcase(key) }"#,
51                result: Ok(r#"{ "A": 1, "B": 2 }"#),
52            },
53            Example {
54                title: "recursively map object keys",
55                source: r#"map_keys({ "a": 1, "b": [{ "c": 2 }, { "d": 3 }], "e": { "f": 4 } }, recursive: true) -> |key| { upcase(key) }"#,
56                result: Ok(r#"{ "A": 1, "B": [{ "C": 2 }, { "D": 3 }], "E": { "F": 4 } }"#),
57            },
58            Example {
59                title: "map nested object keys",
60                source: r#"map_keys({ "a": 1, "b": { "c": 2, "d": 3, "e": { "f": 4 } } }.b) -> |key| { upcase(key) }"#,
61                result: Ok(r#"{ "C": 2, "D": 3, "E": { "f": 4 } }"#),
62            },
63        ]
64    }
65
66    fn compile(
67        &self,
68        _state: &state::TypeState,
69        _ctx: &mut FunctionCompileContext,
70        arguments: ArgumentList,
71    ) -> Compiled {
72        let value = arguments.required("value");
73        let recursive = arguments.optional("recursive");
74        let closure = arguments.required_closure()?;
75
76        Ok(MapKeysFn {
77            value,
78            recursive,
79            closure,
80        }
81        .as_expr())
82    }
83
84    fn closure(&self) -> Option<closure::Definition> {
85        use closure::{Definition, Input, Output, Variable, VariableKind};
86
87        Some(Definition {
88            inputs: vec![Input {
89                parameter_keyword: "value",
90                kind: Kind::object(Collection::any()),
91                variables: vec![Variable {
92                    kind: VariableKind::Exact(Kind::bytes()),
93                }],
94                output: Output::Kind(Kind::bytes()),
95                example: Example {
96                    title: "map object keys",
97                    source: r#"map_keys({ "one" : 1, "two": 2 }) -> |key| { upcase(key) }"#,
98                    result: Ok(r#"{ "ONE": 1, "TWO": 2 }"#),
99                },
100            }],
101            is_iterator: true,
102        })
103    }
104}
105
106#[derive(Debug, Clone)]
107struct MapKeysFn {
108    value: Box<dyn Expression>,
109    recursive: Option<Box<dyn Expression>>,
110    closure: Closure,
111}
112
113impl FunctionExpression for MapKeysFn {
114    fn resolve(&self, ctx: &mut Context) -> ExpressionResult<Value> {
115        let recursive = match &self.recursive {
116            None => false,
117            Some(expr) => expr.resolve(ctx)?.try_boolean()?,
118        };
119
120        let value = self.value.resolve(ctx)?;
121        let Closure {
122            variables,
123            block,
124            block_type_def: _,
125        } = &self.closure;
126        let runner = closure::Runner::new(variables, |ctx| block.resolve(ctx));
127
128        map_keys(value, recursive, ctx, &runner)
129    }
130
131    fn type_def(&self, ctx: &state::TypeState) -> TypeDef {
132        self.value.type_def(ctx)
133    }
134}