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
149 match input {
150 PipelineData::Value(Value::Error { error, .. }, ..) => Err(*error),
152 PipelineData::Value(mut value, metadata) => {
153 if let Value::Closure { val, .. } = replacement {
154 match (cell_path.members.first(), &mut value) {
155 (Some(PathMember::String { .. }), Value::List { vals, .. }) => {
156 let mut closure = ClosureEval::new(engine_state, stack, *val);
157 for val in vals {
158 insert_value_by_closure(
159 val,
160 &mut closure,
161 head,
162 &cell_path.members,
163 false,
164 )?;
165 }
166 }
167 (first, _) => {
168 insert_single_value_by_closure(
169 &mut value,
170 ClosureEvalOnce::new(engine_state, stack, *val),
171 head,
172 &cell_path.members,
173 matches!(first, Some(PathMember::Int { .. })),
174 )?;
175 }
176 }
177 } else {
178 value.insert_data_at_cell_path(&cell_path.members, replacement, head)?;
179 }
180 Ok(value.into_pipeline_data_with_metadata(metadata))
181 }
182 PipelineData::ListStream(stream, metadata) => {
183 if let Some((
184 &PathMember::Int {
185 val,
186 span: path_span,
187 ..
188 },
189 path,
190 )) = cell_path.members.split_first()
191 {
192 let mut stream = stream.into_iter();
193 let mut pre_elems = vec![];
194
195 for idx in 0..val {
196 if let Some(v) = stream.next() {
197 pre_elems.push(v);
198 } else {
199 return Err(ShellError::InsertAfterNextFreeIndex {
200 available_idx: idx,
201 span: path_span,
202 });
203 }
204 }
205
206 if path.is_empty() {
207 if let Value::Closure { val, .. } = replacement {
208 let value = stream.next();
209 let end_of_stream = value.is_none();
210 let value = value.unwrap_or(Value::nothing(head));
211 let new_value = ClosureEvalOnce::new(engine_state, stack, *val)
212 .run_with_value(value.clone())?
213 .into_value(head)?;
214
215 pre_elems.push(new_value);
216 if !end_of_stream {
217 pre_elems.push(value);
218 }
219 } else {
220 pre_elems.push(replacement);
221 }
222 } else if let Some(mut value) = stream.next() {
223 if let Value::Closure { val, .. } = replacement {
224 insert_single_value_by_closure(
225 &mut value,
226 ClosureEvalOnce::new(engine_state, stack, *val),
227 head,
228 path,
229 true,
230 )?;
231 } else {
232 value.insert_data_at_cell_path(path, replacement, head)?;
233 }
234 pre_elems.push(value)
235 } else {
236 return Err(ShellError::AccessBeyondEnd {
237 max_idx: pre_elems.len() - 1,
238 span: path_span,
239 });
240 }
241
242 Ok(pre_elems
243 .into_iter()
244 .chain(stream)
245 .into_pipeline_data_with_metadata(
246 head,
247 engine_state.signals().clone(),
248 metadata,
249 ))
250 } else if let Value::Closure { val, .. } = replacement {
251 let mut closure = ClosureEval::new(engine_state, stack, *val);
252 let stream = stream.map(move |mut value| {
253 let err = insert_value_by_closure(
254 &mut value,
255 &mut closure,
256 head,
257 &cell_path.members,
258 false,
259 );
260
261 if let Err(e) = err {
262 Value::error(e, head)
263 } else {
264 value
265 }
266 });
267 Ok(PipelineData::list_stream(stream, metadata))
268 } else {
269 let stream = stream.map(move |mut value| {
270 if let Err(e) = value.insert_data_at_cell_path(
271 &cell_path.members,
272 replacement.clone(),
273 head,
274 ) {
275 Value::error(e, head)
276 } else {
277 value
278 }
279 });
280
281 Ok(PipelineData::list_stream(stream, metadata))
282 }
283 }
284 PipelineData::Empty => Err(ShellError::IncompatiblePathAccess {
285 type_name: "empty pipeline".to_string(),
286 span: head,
287 }),
288 PipelineData::ByteStream(stream, ..) => Err(ShellError::IncompatiblePathAccess {
289 type_name: stream.type_().describe().into(),
290 span: head,
291 }),
292 }
293}
294
295fn insert_value_by_closure(
296 value: &mut Value,
297 closure: &mut ClosureEval,
298 span: Span,
299 cell_path: &[PathMember],
300 first_path_member_int: bool,
301) -> Result<(), ShellError> {
302 let value_at_path = if first_path_member_int {
303 value
304 .follow_cell_path(cell_path)
305 .map(Cow::into_owned)
306 .unwrap_or(Value::nothing(span))
307 } else {
308 value.clone()
309 };
310
311 let new_value = closure.run_with_value(value_at_path)?.into_value(span)?;
312 value.insert_data_at_cell_path(cell_path, new_value, span)
313}
314
315fn insert_single_value_by_closure(
316 value: &mut Value,
317 closure: ClosureEvalOnce,
318 span: Span,
319 cell_path: &[PathMember],
320 first_path_member_int: bool,
321) -> Result<(), ShellError> {
322 let value_at_path = if first_path_member_int {
323 value
324 .follow_cell_path(cell_path)
325 .map(Cow::into_owned)
326 .unwrap_or(Value::nothing(span))
327 } else {
328 value.clone()
329 };
330
331 let new_value = closure.run_with_value(value_at_path)?.into_value(span)?;
332 value.insert_data_at_cell_path(cell_path, new_value, span)
333}
334
335#[cfg(test)]
336mod test {
337 use super::*;
338
339 #[test]
340 fn test_examples() {
341 use crate::test_examples;
342
343 test_examples(Insert {})
344 }
345}