Add -c flag to update cells subcommand (#4039)

* Add `-c` flag to `update cells` subcommand

* Fix lints

crates/nu-command/src/commands/filters/update_cells.rs

@@ -7,6 +7,8 @@ use nu_protocol::{
     hir::{CapturedBlock, ExternalRedirection},
     Signature, SyntaxShape, TaggedDictBuilder, UntaggedValue, Value,
 };
+use std::collections::HashSet;
+use std::iter::FromIterator;
 
 pub struct SubCommand;
 
@@ -16,11 +18,18 @@ impl WholeStreamCommand for SubCommand {
     }
 
     fn signature(&self) -> Signature {
-        Signature::build("update cells").required(
+        Signature::build("update cells")
-            "block",
+            .required(
-            SyntaxShape::Block,
+                "block",
-            "the block to run an update for each cell",
+                SyntaxShape::Block,
-        )
+                "the block to run an update for each cell",
+            )
+            .named(
+                "columns",
+                SyntaxShape::Table,
+                "list of columns to update",
+                Some('c'),
+            )
     }
 
     fn usage(&self) -> &str {
@@ -32,9 +41,10 @@ impl WholeStreamCommand for SubCommand {
     }
 
     fn examples(&self) -> Vec<Example> {
-        vec![Example {
+        vec![
-            description: "Update the zero value cells to empty strings.",
+            Example {
-            example: r#"[
+                description: "Update the zero value cells to empty strings.",
+                example: r#"[
   [2021-04-16, 2021-06-10, 2021-09-18, 2021-10-15, 2021-11-16, 2021-11-17, 2021-11-18];
   [        37,          0,          0,          0,         37,          0,          0]
 ] | update cells {|value|
@@ -44,17 +54,41 @@ impl WholeStreamCommand for SubCommand {
         $value
       }
 }"#,
-            result: Some(vec![UntaggedValue::row(indexmap! {
+                result: Some(vec![UntaggedValue::row(indexmap! {
-                    "2021-04-16".to_string() => UntaggedValue::int(37).into(),
+                        "2021-04-16".to_string() => UntaggedValue::int(37).into(),
-                    "2021-06-10".to_string() => Value::from(""),
+                        "2021-06-10".to_string() => Value::from(""),
-                    "2021-09-18".to_string() => Value::from(""),
+                        "2021-09-18".to_string() => Value::from(""),
-                    "2021-10-15".to_string() => Value::from(""),
+                        "2021-10-15".to_string() => Value::from(""),
-                    "2021-11-16".to_string() => UntaggedValue::int(37).into(),
+                        "2021-11-16".to_string() => UntaggedValue::int(37).into(),
-                    "2021-11-17".to_string() => Value::from(""),
+                        "2021-11-17".to_string() => Value::from(""),
-                    "2021-11-18".to_string() => Value::from(""),
+                        "2021-11-18".to_string() => Value::from(""),
-            })
+                })
-            .into()]),
+                .into()]),
-        }]
+            },
+            Example {
+                description: "Update the zero value cells to empty strings in 2 last columns.",
+                example: r#"[
+    [2021-04-16, 2021-06-10, 2021-09-18, 2021-10-15, 2021-11-16, 2021-11-17, 2021-11-18];
+    [        37,          0,          0,          0,         37,          0,          0]
+] | update cells -c ["2021-11-18", "2021-11-17"] {|value|
+        if ($value | into int) == 0 {
+        ""
+        } {
+        $value
+        }
+}"#,
+                result: Some(vec![UntaggedValue::row(indexmap! {
+                        "2021-04-16".to_string() => UntaggedValue::int(37).into(),
+                        "2021-06-10".to_string() => UntaggedValue::int(0).into(),
+                        "2021-09-18".to_string() => UntaggedValue::int(0).into(),
+                        "2021-10-15".to_string() => UntaggedValue::int(0).into(),
+                        "2021-11-16".to_string() => UntaggedValue::int(37).into(),
+                        "2021-11-17".to_string() => Value::from(""),
+                        "2021-11-18".to_string() => Value::from(""),
+                })
+                .into()]),
+            },
+        ]
     }
 }
 
@@ -65,6 +99,11 @@ fn update_cells(args: CommandArgs) -> Result<OutputStream, ShellError> {
     let block: CapturedBlock = args.req(0)?;
     let block = Arc::new(block);
 
+    let columns = args
+        .get_flag("columns")?
+        .map(|x: Value| HashSet::from_iter(x.table_entries().map(|val| val.convert_to_string())));
+    let columns = Arc::new(columns);
+
     Ok(args
         .input
         .flat_map(move |input| {
@@ -72,7 +111,13 @@ fn update_cells(args: CommandArgs) -> Result<OutputStream, ShellError> {
             let context = context.clone();
 
             if input.is_row() {
-                OutputStream::one(process_cells(block, context, input, external_redirection))
+                OutputStream::one(process_cells(
+                    block,
+                    columns.clone(),
+                    context,
+                    input,
+                    external_redirection,
+                ))
             } else {
                 match process_input(block, context, input, external_redirection) {
                     Ok(s) => s,
@@ -123,12 +168,20 @@ pub fn process_input(
 
 pub fn process_cells(
     captured_block: Arc<CapturedBlock>,
+    columns: Arc<Option<HashSet<String>>>,
     context: Arc<EvaluationContext>,
     input: Value,
     external_redirection: ExternalRedirection,
 ) -> Value {
     TaggedDictBuilder::build(input.tag(), |row| {
         input.row_entries().for_each(|(column, cell_value)| {
+            match &*columns {
+                Some(col) if !col.contains(column) => {
+                    row.insert_value(column, cell_value.clone());
+                    return;
+                }
+                _ => {}
+            };
             let cell_processed = process_input(
                 captured_block.clone(),
                 context.clone(),