Skip to main content

nu_command/filters/
insert.rs

1use std::borrow::Cow;
2
3use nu_engine::{ClosureEval, ClosureEvalOnce, command_prelude::*};
4use nu_protocol::ast::PathMember;
5
6#[derive(Clone)]
7pub struct Insert;
8
9impl Command for Insert {
10    fn name(&self) -> &str {
11        "insert"
12    }
13
14    fn signature(&self) -> Signature {
15        Signature::build("insert")
16            .input_output_types(vec![
17                (Type::record(), Type::record()),
18                (Type::table(), Type::table()),
19                (
20                    Type::List(Box::new(Type::Any)),
21                    Type::List(Box::new(Type::Any)),
22                ),
23            ])
24            .required(
25                "field",
26                SyntaxShape::CellPath,
27                "The name of the column to insert.",
28            )
29            .required(
30                "new value",
31                SyntaxShape::Any,
32                "The new value to give the cell(s).",
33            )
34            .allow_variants_without_examples(true)
35            .category(Category::Filters)
36    }
37
38    fn description(&self) -> &str {
39        "Insert a new column, using an expression or closure to create each row's values."
40    }
41
42    fn extra_description(&self) -> &str {
43        "When inserting a column, the closure will be run for each row, and the current row will be passed as the first argument.
44When inserting into a specific index, the closure will instead get the current value at the index or null if inserting at the end of a list/table."
45    }
46
47    fn search_terms(&self) -> Vec<&str> {
48        vec!["add"]
49    }
50
51    fn run(
52        &self,
53        engine_state: &EngineState,
54        stack: &mut Stack,
55        call: &Call,
56        input: PipelineData,
57    ) -> Result<PipelineData, ShellError> {
58        insert(engine_state, stack, call, input)
59    }
60
61    fn examples(&self) -> Vec<Example<'_>> {
62        vec![
63            Example {
64                description: "Insert a new entry into a single record",
65                example: "{'name': 'nu', 'stars': 5} | insert alias 'Nushell'",
66                result: Some(Value::test_record(record! {
67                    "name" => Value::test_string("nu"),
68                    "stars" => Value::test_int(5),
69                    "alias" => Value::test_string("Nushell"),
70                })),
71            },
72            Example {
73                description: "Insert a new column into a table, populating all rows",
74                example: "[[project, lang]; ['Nushell', 'Rust']] | insert type 'shell'",
75                result: Some(Value::test_list(vec![Value::test_record(record! {
76                    "project" => Value::test_string("Nushell"),
77                    "lang" => Value::test_string("Rust"),
78                    "type" => Value::test_string("shell"),
79                })])),
80            },
81            Example {
82                description: "Insert a new column with values computed based off the other columns",
83                example: "[[foo]; [7] [8] [9]] | insert bar {|row| $row.foo * 2 }",
84                result: Some(Value::test_list(vec![
85                    Value::test_record(record! {
86                        "foo" => Value::test_int(7),
87                        "bar" => Value::test_int(14),
88                    }),
89                    Value::test_record(record! {
90                        "foo" => Value::test_int(8),
91                        "bar" => Value::test_int(16),
92                    }),
93                    Value::test_record(record! {
94                        "foo" => Value::test_int(9),
95                        "bar" => Value::test_int(18),
96                    }),
97                ])),
98            },
99            Example {
100                description: "Insert a new value into a list at an index",
101                example: "[1 2 4] | insert 2 3",
102                result: Some(Value::test_list(vec![
103                    Value::test_int(1),
104                    Value::test_int(2),
105                    Value::test_int(3),
106                    Value::test_int(4),
107                ])),
108            },
109            Example {
110                description: "Insert a new value at the end of a list",
111                example: "[1 2 3] | insert 3 4",
112                result: Some(Value::test_list(vec![
113                    Value::test_int(1),
114                    Value::test_int(2),
115                    Value::test_int(3),
116                    Value::test_int(4),
117                ])),
118            },
119            Example {
120                description: "Insert into a nested path, creating new values as needed",
121                example: "[{} {a: [{}]}] | insert a.0.b \"value\"",
122                result: Some(Value::test_list(vec![
123                    Value::test_record(record!(
124                        "a" => Value::test_list(vec![Value::test_record(record!(
125                            "b" => Value::test_string("value"),
126                        ))]),
127                    )),
128                    Value::test_record(record!(
129                        "a" => Value::test_list(vec![Value::test_record(record!(
130                            "b" => Value::test_string("value"),
131                        ))]),
132                    )),
133                ])),
134            },
135        ]
136    }
137}
138
139fn insert(
140    engine_state: &EngineState,
141    stack: &mut Stack,
142    call: &Call,
143    input: PipelineData,
144) -> Result<PipelineData, ShellError> {
145    let head = call.head;
146    let cell_path: CellPath = call.req(engine_state, stack, 0)?;
147    let replacement: Value = call.req(engine_state, stack, 1)?;
148    let input = match input.try_into_stream(engine_state) {
149        Ok(input) | Err(input) => input,
150    };
151
152    match input {
153        // Propagate errors in the pipeline
154        PipelineData::Value(Value::Error { error, .. }, ..) => Err(*error),
155        PipelineData::Value(mut value, metadata) => {
156            if let Value::Closure { val, .. } = replacement {
157                match (cell_path.members.first(), &mut value) {
158                    (Some(PathMember::String { .. }), Value::List { vals, .. }) => {
159                        let mut closure = ClosureEval::new(engine_state, stack, *val);
160                        for val in vals {
161                            insert_value_by_closure(
162                                val,
163                                &mut closure,
164                                head,
165                                &cell_path.members,
166                                false,
167                            )?;
168                        }
169                    }
170                    (first, _) => {
171                        insert_single_value_by_closure(
172                            &mut value,
173                            ClosureEvalOnce::new(engine_state, stack, *val),
174                            head,
175                            &cell_path.members,
176                            matches!(first, Some(PathMember::Int { .. })),
177                        )?;
178                    }
179                }
180            } else {
181                value.insert_data_at_cell_path(&cell_path.members, replacement, head)?;
182            }
183            Ok(value.into_pipeline_data_with_metadata(metadata))
184        }
185        PipelineData::ListStream(stream, metadata) => {
186            if let Some((
187                &PathMember::Int {
188                    val,
189                    span: path_span,
190                    ..
191                },
192                path,
193            )) = cell_path.members.split_first()
194            {
195                let mut stream = stream.into_iter();
196                let mut pre_elems = vec![];
197
198                for idx in 0..val {
199                    if let Some(v) = stream.next() {
200                        pre_elems.push(v);
201                    } else {
202                        return Err(ShellError::InsertAfterNextFreeIndex {
203                            available_idx: idx,
204                            span: path_span,
205                        });
206                    }
207                }
208
209                if path.is_empty() {
210                    if let Value::Closure { val, .. } = replacement {
211                        let value = stream.next();
212                        let end_of_stream = value.is_none();
213                        let value = value.unwrap_or(Value::nothing(head));
214                        let new_value = ClosureEvalOnce::new(engine_state, stack, *val)
215                            .run_with_value(value.clone())?
216                            .into_value(head)?;
217
218                        pre_elems.push(new_value);
219                        if !end_of_stream {
220                            pre_elems.push(value);
221                        }
222                    } else {
223                        pre_elems.push(replacement);
224                    }
225                } else if let Some(mut value) = stream.next() {
226                    if let Value::Closure { val, .. } = replacement {
227                        insert_single_value_by_closure(
228                            &mut value,
229                            ClosureEvalOnce::new(engine_state, stack, *val),
230                            head,
231                            path,
232                            true,
233                        )?;
234                    } else {
235                        value.insert_data_at_cell_path(path, replacement, head)?;
236                    }
237                    pre_elems.push(value)
238                } else {
239                    return Err(ShellError::AccessBeyondEnd {
240                        max_idx: pre_elems.len() - 1,
241                        span: path_span,
242                    });
243                }
244
245                Ok(pre_elems
246                    .into_iter()
247                    .chain(stream)
248                    .into_pipeline_data_with_metadata(
249                        head,
250                        engine_state.signals().clone(),
251                        metadata,
252                    ))
253            } else if let Value::Closure { val, .. } = replacement {
254                let mut closure = ClosureEval::new(engine_state, stack, *val);
255                let stream = stream.map(move |mut value| {
256                    let err = insert_value_by_closure(
257                        &mut value,
258                        &mut closure,
259                        head,
260                        &cell_path.members,
261                        false,
262                    );
263
264                    if let Err(e) = err {
265                        Value::error(e, head)
266                    } else {
267                        value
268                    }
269                });
270                Ok(PipelineData::list_stream(stream, metadata))
271            } else {
272                let stream = stream.map(move |mut value| {
273                    if let Err(e) = value.insert_data_at_cell_path(
274                        &cell_path.members,
275                        replacement.clone(),
276                        head,
277                    ) {
278                        Value::error(e, head)
279                    } else {
280                        value
281                    }
282                });
283
284                Ok(PipelineData::list_stream(stream, metadata))
285            }
286        }
287        PipelineData::Empty => Err(ShellError::IncompatiblePathAccess {
288            type_name: "empty pipeline".to_string(),
289            span: head,
290        }),
291        PipelineData::ByteStream(stream, ..) => Err(ShellError::IncompatiblePathAccess {
292            type_name: stream.type_().describe().into(),
293            span: head,
294        }),
295    }
296}
297
298fn insert_value_by_closure(
299    value: &mut Value,
300    closure: &mut ClosureEval,
301    span: Span,
302    cell_path: &[PathMember],
303    first_path_member_int: bool,
304) -> Result<(), ShellError> {
305    let value_at_path = if first_path_member_int {
306        value
307            .follow_cell_path(cell_path)
308            .map(Cow::into_owned)
309            .unwrap_or(Value::nothing(span))
310    } else {
311        value.clone()
312    };
313
314    let new_value = closure.run_with_value(value_at_path)?.into_value(span)?;
315    value.insert_data_at_cell_path(cell_path, new_value, span)
316}
317
318fn insert_single_value_by_closure(
319    value: &mut Value,
320    closure: ClosureEvalOnce,
321    span: Span,
322    cell_path: &[PathMember],
323    first_path_member_int: bool,
324) -> Result<(), ShellError> {
325    let value_at_path = if first_path_member_int {
326        value
327            .follow_cell_path(cell_path)
328            .map(Cow::into_owned)
329            .unwrap_or(Value::nothing(span))
330    } else {
331        value.clone()
332    };
333
334    let new_value = closure.run_with_value(value_at_path)?.into_value(span)?;
335    value.insert_data_at_cell_path(cell_path, new_value, span)
336}
337
338#[cfg(test)]
339mod test {
340    use super::*;
341
342    #[test]
343    fn test_examples() {
344        use crate::test_examples;
345
346        test_examples(Insert {})
347    }
348}