1use nu_engine::{ClosureEval, ClosureEvalOnce, command_prelude::*};
2use nu_protocol::ast::PathMember;
3
4#[derive(Clone)]
5pub struct Update;
6
7impl Command for Update {
8 fn name(&self) -> &str {
9 "update"
10 }
11
12 fn signature(&self) -> Signature {
13 Signature::build("update")
14 .input_output_types(vec![
15 (Type::record(), Type::record()),
16 (Type::table(), Type::table()),
17 (
18 Type::List(Box::new(Type::Any)),
19 Type::List(Box::new(Type::Any)),
20 ),
21 ])
22 .required(
23 "field",
24 SyntaxShape::CellPath,
25 "The name of the column to update.",
26 )
27 .required(
28 "replacement value",
29 SyntaxShape::Any,
30 "The new value to give the cell(s), or a closure to create the value.",
31 )
32 .allow_variants_without_examples(true)
33 .category(Category::Filters)
34 }
35
36 fn description(&self) -> &str {
37 "Update an existing column to have a new value."
38 }
39
40 fn extra_description(&self) -> &str {
41 "When updating a column, the closure will be run for each row, and the current row will be passed as the first argument. \
42Referencing `$in` inside the closure will provide the value at the column for the current row.
43
44When updating a specific index, the closure will instead be run once. The first argument to the closure and the `$in` value will both be the current value at the index."
45 }
46
47 fn run(
48 &self,
49 engine_state: &EngineState,
50 stack: &mut Stack,
51 call: &Call,
52 input: PipelineData,
53 ) -> Result<PipelineData, ShellError> {
54 update(engine_state, stack, call, input)
55 }
56
57 fn examples(&self) -> Vec<Example> {
58 vec![
59 Example {
60 description: "Update a column value",
61 example: "{'name': 'nu', 'stars': 5} | update name 'Nushell'",
62 result: Some(Value::test_record(record! {
63 "name" => Value::test_string("Nushell"),
64 "stars" => Value::test_int(5),
65 })),
66 },
67 Example {
68 description: "Use a closure to alter each value in the 'authors' column to a single string",
69 example: "[[project, authors]; ['nu', ['Andrés', 'JT', 'Yehuda']]] | update authors {|row| $row.authors | str join ',' }",
70 result: Some(Value::test_list(vec![Value::test_record(record! {
71 "project" => Value::test_string("nu"),
72 "authors" => Value::test_string("Andrés,JT,Yehuda"),
73 })])),
74 },
75 Example {
76 description: "Implicitly use the `$in` value in a closure to update 'authors'",
77 example: "[[project, authors]; ['nu', ['Andrés', 'JT', 'Yehuda']]] | update authors { str join ',' }",
78 result: Some(Value::test_list(vec![Value::test_record(record! {
79 "project" => Value::test_string("nu"),
80 "authors" => Value::test_string("Andrés,JT,Yehuda"),
81 })])),
82 },
83 Example {
84 description: "Update a value at an index in a list",
85 example: "[1 2 3] | update 1 4",
86 result: Some(Value::test_list(vec![
87 Value::test_int(1),
88 Value::test_int(4),
89 Value::test_int(3),
90 ])),
91 },
92 Example {
93 description: "Use a closure to compute a new value at an index",
94 example: "[1 2 3] | update 1 {|i| $i + 2 }",
95 result: Some(Value::test_list(vec![
96 Value::test_int(1),
97 Value::test_int(4),
98 Value::test_int(3),
99 ])),
100 },
101 ]
102 }
103}
104
105fn update(
106 engine_state: &EngineState,
107 stack: &mut Stack,
108 call: &Call,
109 input: PipelineData,
110) -> Result<PipelineData, ShellError> {
111 let head = call.head;
112 let cell_path: CellPath = call.req(engine_state, stack, 0)?;
113 let replacement: Value = call.req(engine_state, stack, 1)?;
114
115 match input {
116 PipelineData::Value(mut value, metadata) => {
117 if let Value::Closure { val, .. } = replacement {
118 match (cell_path.members.first(), &mut value) {
119 (Some(PathMember::String { .. }), Value::List { vals, .. }) => {
120 let mut closure = ClosureEval::new(engine_state, stack, *val);
121 for val in vals {
122 update_value_by_closure(
123 val,
124 &mut closure,
125 head,
126 &cell_path.members,
127 false,
128 )?;
129 }
130 }
131 (first, _) => {
132 update_single_value_by_closure(
133 &mut value,
134 ClosureEvalOnce::new(engine_state, stack, *val),
135 head,
136 &cell_path.members,
137 matches!(first, Some(PathMember::Int { .. })),
138 )?;
139 }
140 }
141 } else {
142 value.update_data_at_cell_path(&cell_path.members, replacement)?;
143 }
144 Ok(value.into_pipeline_data_with_metadata(metadata))
145 }
146 PipelineData::ListStream(stream, metadata) => {
147 if let Some((
148 &PathMember::Int {
149 val,
150 span: path_span,
151 ..
152 },
153 path,
154 )) = cell_path.members.split_first()
155 {
156 let mut stream = stream.into_iter();
157 let mut pre_elems = vec![];
158
159 for idx in 0..=val {
160 if let Some(v) = stream.next() {
161 pre_elems.push(v);
162 } else if idx == 0 {
163 return Err(ShellError::AccessEmptyContent { span: path_span });
164 } else {
165 return Err(ShellError::AccessBeyondEnd {
166 max_idx: idx - 1,
167 span: path_span,
168 });
169 }
170 }
171
172 let value = pre_elems.last_mut().expect("one element");
174
175 if let Value::Closure { val, .. } = replacement {
176 update_single_value_by_closure(
177 value,
178 ClosureEvalOnce::new(engine_state, stack, *val),
179 head,
180 path,
181 true,
182 )?;
183 } else {
184 value.update_data_at_cell_path(path, replacement)?;
185 }
186
187 Ok(pre_elems
188 .into_iter()
189 .chain(stream)
190 .into_pipeline_data_with_metadata(
191 head,
192 engine_state.signals().clone(),
193 metadata,
194 ))
195 } else if let Value::Closure { val, .. } = replacement {
196 let mut closure = ClosureEval::new(engine_state, stack, *val);
197 let stream = stream.map(move |mut value| {
198 let err = update_value_by_closure(
199 &mut value,
200 &mut closure,
201 head,
202 &cell_path.members,
203 false,
204 );
205
206 if let Err(e) = err {
207 Value::error(e, head)
208 } else {
209 value
210 }
211 });
212
213 Ok(PipelineData::list_stream(stream, metadata))
214 } else {
215 let stream = stream.map(move |mut value| {
216 if let Err(e) =
217 value.update_data_at_cell_path(&cell_path.members, replacement.clone())
218 {
219 Value::error(e, head)
220 } else {
221 value
222 }
223 });
224
225 Ok(PipelineData::list_stream(stream, metadata))
226 }
227 }
228 PipelineData::Empty => Err(ShellError::IncompatiblePathAccess {
229 type_name: "empty pipeline".to_string(),
230 span: head,
231 }),
232 PipelineData::ByteStream(stream, ..) => Err(ShellError::IncompatiblePathAccess {
233 type_name: stream.type_().describe().into(),
234 span: head,
235 }),
236 }
237}
238
239fn update_value_by_closure(
240 value: &mut Value,
241 closure: &mut ClosureEval,
242 span: Span,
243 cell_path: &[PathMember],
244 first_path_member_int: bool,
245) -> Result<(), ShellError> {
246 let value_at_path = value.follow_cell_path(cell_path)?;
247
248 let arg = if first_path_member_int {
249 value_at_path.as_ref()
250 } else {
251 &*value
252 };
253
254 let new_value = closure
255 .add_arg(arg.clone())
256 .run_with_input(value_at_path.into_owned().into_pipeline_data())?
257 .into_value(span)?;
258
259 value.update_data_at_cell_path(cell_path, new_value)
260}
261
262fn update_single_value_by_closure(
263 value: &mut Value,
264 closure: ClosureEvalOnce,
265 span: Span,
266 cell_path: &[PathMember],
267 first_path_member_int: bool,
268) -> Result<(), ShellError> {
269 let value_at_path = value.follow_cell_path(cell_path)?;
270
271 let arg = if first_path_member_int {
272 value_at_path.as_ref()
273 } else {
274 &*value
275 };
276
277 let new_value = closure
278 .add_arg(arg.clone())
279 .run_with_input(value_at_path.into_owned().into_pipeline_data())?
280 .into_value(span)?;
281
282 value.update_data_at_cell_path(cell_path, new_value)
283}
284
285#[cfg(test)]
286mod test {
287 use super::*;
288
289 #[test]
290 fn test_examples() {
291 use crate::test_examples;
292
293 test_examples(Update {})
294 }
295}