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}