vrl/stdlib/
map_keys.rs

1use crate::compiler::prelude::*;
2use std::sync::LazyLock;
3
4static DEFAULT_RECURSIVE: LazyLock<Value> = LazyLock::new(|| Value::Boolean(false));
5
6static PARAMETERS: LazyLock<Vec<Parameter>> = LazyLock::new(|| {
7    vec![
8        Parameter::required("value", kind::OBJECT, "The object to iterate."),
9        Parameter::optional(
10            "recursive",
11            kind::BOOLEAN,
12            "Whether to recursively iterate the collection.",
13        )
14        .default(&DEFAULT_RECURSIVE),
15    ]
16});
17
18fn map_keys<T>(
19    value: Value,
20    recursive: bool,
21    ctx: &mut Context,
22    runner: &closure::Runner<T>,
23) -> Resolved
24where
25    T: Fn(&mut Context) -> Resolved,
26{
27    let mut iter = value.into_iter(recursive);
28
29    for item in iter.by_ref() {
30        if let IterItem::KeyValue(key, _) = item {
31            runner.map_key(ctx, key)?;
32        }
33    }
34
35    Ok(iter.into())
36}
37
38#[derive(Clone, Copy, Debug)]
39pub struct MapKeys;
40
41impl Function for MapKeys {
42    fn identifier(&self) -> &'static str {
43        "map_keys"
44    }
45
46    fn usage(&self) -> &'static str {
47        indoc! {"
48            Map the keys within an object.
49
50            If `recursive` is enabled, the function iterates into nested
51            objects, using the following rules:
52
53            1. Iteration starts at the root.
54            2. For every nested object type:
55               - First return the key of the object type itself.
56               - Then recurse into the object, and loop back to item (1)
57                 in this list.
58               - Any mutation done on a nested object *before* recursing into
59                 it, are preserved.
60            3. For every nested array type:
61               - First return the key of the array type itself.
62               - Then find all objects within the array, and apply item (2)
63                 to each individual object.
64
65            The above rules mean that `map_keys` with
66            `recursive` enabled finds *all* keys in the target,
67            regardless of whether nested objects are nested inside arrays.
68
69            The function uses the function closure syntax to allow reading
70            the key for each item in the object.
71
72            The same scoping rules apply to closure blocks as they do for
73            regular blocks. This means that any variable defined in parent scopes
74            is accessible, and mutations to those variables are preserved,
75            but any new variables instantiated in the closure block are
76            unavailable outside of the block.
77
78            See the examples below to learn about the closure syntax.
79        "}
80    }
81
82    fn category(&self) -> &'static str {
83        Category::Enumerate.as_ref()
84    }
85
86    fn return_kind(&self) -> u16 {
87        kind::OBJECT
88    }
89
90    fn parameters(&self) -> &'static [Parameter] {
91        PARAMETERS.as_slice()
92    }
93
94    fn examples(&self) -> &'static [Example] {
95        &[
96            example! {
97                title: "Upcase keys",
98                source: indoc! {r#"
99                    . = {
100                        "foo": "foo",
101                        "bar": "bar",
102                        "baz": {"nested key": "val"}
103                    }
104                    map_keys(.) -> |key| { upcase(key) }
105                "#},
106                result: Ok(r#"{ "FOO": "foo", "BAR": "bar", "BAZ": {"nested key": "val"} }"#),
107            },
108            example! {
109                title: "De-dot keys",
110                source: indoc! {r#"
111                    . = {
112                        "labels": {
113                            "app.kubernetes.io/name": "mysql"
114                        }
115                    }
116                    map_keys(., recursive: true) -> |key| { replace(key, ".", "_") }
117                "#},
118                result: Ok(r#"{ "labels": { "app_kubernetes_io/name": "mysql" } }"#),
119            },
120            example! {
121                title: "Recursively map object keys",
122                source: indoc! {r#"
123                    val = {
124                        "a": 1,
125                        "b": [{ "c": 2 }, { "d": 3 }],
126                        "e": { "f": 4 }
127                    }
128                    map_keys(val, recursive: true) -> |key| { upcase(key) }
129                "#},
130                result: Ok(r#"{ "A": 1, "B": [{ "C": 2 }, { "D": 3 }], "E": { "F": 4 } }"#),
131            },
132        ]
133    }
134
135    fn compile(
136        &self,
137        _state: &state::TypeState,
138        _ctx: &mut FunctionCompileContext,
139        arguments: ArgumentList,
140    ) -> Compiled {
141        let value = arguments.required("value");
142        let recursive = arguments.optional("recursive");
143        let closure = arguments.required_closure()?;
144
145        Ok(MapKeysFn {
146            value,
147            recursive,
148            closure,
149        }
150        .as_expr())
151    }
152
153    fn closure(&self) -> Option<closure::Definition> {
154        use closure::{Definition, Input, Output, Variable, VariableKind};
155
156        Some(Definition {
157            inputs: vec![Input {
158                parameter_keyword: "value",
159                kind: Kind::object(Collection::any()),
160                variables: vec![Variable {
161                    kind: VariableKind::Exact(Kind::bytes()),
162                }],
163                output: Output::Kind(Kind::bytes()),
164                example: example! {
165                    title: "map object keys",
166                    source: r#"map_keys({ "one" : 1, "two": 2 }) -> |key| { upcase(key) }"#,
167                    result: Ok(r#"{ "ONE": 1, "TWO": 2 }"#),
168                },
169            }],
170            is_iterator: true,
171        })
172    }
173}
174
175#[derive(Debug, Clone)]
176struct MapKeysFn {
177    value: Box<dyn Expression>,
178    recursive: Option<Box<dyn Expression>>,
179    closure: Closure,
180}
181
182impl FunctionExpression for MapKeysFn {
183    fn resolve(&self, ctx: &mut Context) -> ExpressionResult<Value> {
184        let recursive = self
185            .recursive
186            .map_resolve_with_default(ctx, || DEFAULT_RECURSIVE.clone())?
187            .try_boolean()?;
188
189        let value = self.value.resolve(ctx)?;
190        let Closure {
191            variables,
192            block,
193            block_type_def: _,
194        } = &self.closure;
195        let runner = closure::Runner::new(variables, |ctx| block.resolve(ctx));
196
197        map_keys(value, recursive, ctx, &runner)
198    }
199
200    fn type_def(&self, ctx: &state::TypeState) -> TypeDef {
201        self.value.type_def(ctx)
202    }
203}