Port update cells command (#891)

* Port update cells command

Clean up, nicer match statements in UpdateCellsIterator

Return columns flag into HashSet errors

Add FIXME: for update cell behavior on nested lists

* Fix: process cells for Record when no columns are specified

* Fix: address clippy lints for unwrap and into_iter

* Fix: don't step into lists and don't bind $it var

crates/nu-command/src/default_context.rs

@@ -97,6 +97,7 @@ pub fn create_default_context(cwd: impl AsRef<Path>) -> EngineState {
             Transpose,
             Uniq,
             Update,
+            UpdateCells,
             Where,
             Wrap,
             Zip,

crates/nu-command/src/filters/mod.rs

@@ -36,6 +36,7 @@ mod sort_by;
 mod transpose;
 mod uniq;
 mod update;
+mod update_cells;
 mod where_;
 mod wrap;
 mod zip_;
@@ -78,6 +79,7 @@ pub use sort_by::SortBy;
 pub use transpose::Transpose;
 pub use uniq::*;
 pub use update::Update;
+pub use update_cells::UpdateCells;
 pub use where_::Where;
 pub use wrap::Wrap;
 pub use zip_::Zip;

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

@@ -0,0 +1,247 @@
+use nu_engine::{eval_block, CallExt};
+use nu_protocol::ast::{Block, Call};
+use nu_protocol::engine::{CaptureBlock, Command, EngineState, Stack};
+use nu_protocol::{
+    Category, Example, IntoInterruptiblePipelineData, IntoPipelineData, PipelineData,
+    PipelineIterator, ShellError, Signature, Span, SyntaxShape, Value,
+};
+use std::collections::HashSet;
+use std::iter::FromIterator;
+
+#[derive(Clone)]
+pub struct UpdateCells;
+
+impl Command for UpdateCells {
+    fn name(&self) -> &str {
+        "update cells"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build("update cells")
+            .required(
+                "block",
+                SyntaxShape::Block(Some(vec![SyntaxShape::Any])),
+                "the block to run an update for each cell",
+            )
+            .named(
+                "columns",
+                SyntaxShape::Table,
+                "list of columns to update",
+                Some('c'),
+            )
+            .category(Category::Filters)
+    }
+
+    fn usage(&self) -> &str {
+        "Update the table cells."
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![
+            Example {
+                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|
+      if $value == 0 {
+        ""
+      } else {
+        $value
+      }
+}"#,
+                result: Some(Value::List {
+                    vals: vec![Value::Record {
+                        cols: vec![
+                            "2021-04-16".into(),
+                            "2021-06-10".into(),
+                            "2021-09-18".into(),
+                            "2021-10-15".into(),
+                            "2021-11-16".into(),
+                            "2021-11-17".into(),
+                            "2021-11-18".into(),
+                        ],
+                        vals: vec![
+                            Value::test_int(37),
+                            Value::test_string(""),
+                            Value::test_string(""),
+                            Value::test_string(""),
+                            Value::test_int(37),
+                            Value::test_string(""),
+                            Value::test_string(""),
+                        ],
+                        span: Span::test_data(),
+                    }],
+                    span: Span::test_data(),
+                }),
+            },
+            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 == 0 {
+          ""
+        } else {
+          $value
+        }
+}"#,
+                result: Some(Value::List {
+                    vals: vec![Value::Record {
+                        cols: vec![
+                            "2021-04-16".into(),
+                            "2021-06-10".into(),
+                            "2021-09-18".into(),
+                            "2021-10-15".into(),
+                            "2021-11-16".into(),
+                            "2021-11-17".into(),
+                            "2021-11-18".into(),
+                        ],
+                        vals: vec![
+                            Value::test_int(37),
+                            Value::test_int(0),
+                            Value::test_int(0),
+                            Value::test_int(0),
+                            Value::test_int(37),
+                            Value::test_string(""),
+                            Value::test_string(""),
+                        ],
+                        span: Span::test_data(),
+                    }],
+                    span: Span::test_data(),
+                }),
+            },
+        ]
+    }
+
+    fn run(
+        &self,
+        engine_state: &EngineState,
+        stack: &mut Stack,
+        call: &Call,
+        input: PipelineData,
+    ) -> Result<PipelineData, ShellError> {
+        // the block to run on each cell
+        let engine_state = engine_state.clone();
+        let block: CaptureBlock = call.req(&engine_state, stack, 0)?;
+        let mut stack = stack.captures_to_stack(&block.captures);
+        let orig_env_vars = stack.env_vars.clone();
+        let orig_env_hidden = stack.env_hidden.clone();
+
+        let ctrlc = engine_state.ctrlc.clone();
+        let block: Block = engine_state.get_block(block.block_id).clone();
+
+        let span = call.head;
+
+        stack.with_env(&orig_env_vars, &orig_env_hidden);
+
+        // the columns to update
+        let columns: Option<Value> = call.get_flag(&engine_state, &mut stack, "columns")?;
+        let columns: Option<HashSet<String>> = match columns {
+            Some(val) => {
+                let cols = val
+                    .as_list()?
+                    .iter()
+                    .map(|val| val.as_string())
+                    .collect::<Result<Vec<String>, ShellError>>()?;
+                Some(HashSet::from_iter(cols.into_iter()))
+            }
+            None => None,
+        };
+
+        Ok(UpdateCellIterator {
+            input: input.into_iter(),
+            engine_state,
+            stack,
+            block,
+            columns,
+            span,
+        }
+        .into_pipeline_data(ctrlc))
+    }
+}
+
+struct UpdateCellIterator {
+    input: PipelineIterator,
+    columns: Option<HashSet<String>>,
+    engine_state: EngineState,
+    stack: Stack,
+    block: Block,
+    span: Span,
+}
+
+impl Iterator for UpdateCellIterator {
+    type Item = Value;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.input.next() {
+            Some(val) => {
+                if let Some(ref cols) = self.columns {
+                    if !val.columns().iter().any(|c| cols.contains(c)) {
+                        return Some(val);
+                    }
+                }
+
+                match val {
+                    Value::Record { vals, cols, span } => Some(Value::Record {
+                        vals: cols
+                            .iter()
+                            .zip(vals.into_iter())
+                            .map(|(col, val)| match &self.columns {
+                                Some(cols) if !cols.contains(col) => val,
+                                _ => process_cell(
+                                    val,
+                                    &self.engine_state,
+                                    &mut self.stack,
+                                    &self.block,
+                                    span,
+                                ),
+                            })
+                            .collect(),
+                        cols,
+                        span,
+                    }),
+                    val => Some(process_cell(
+                        val,
+                        &self.engine_state,
+                        &mut self.stack,
+                        &self.block,
+                        self.span,
+                    )),
+                }
+            }
+            None => None,
+        }
+    }
+}
+
+fn process_cell(
+    val: Value,
+    engine_state: &EngineState,
+    stack: &mut Stack,
+    block: &Block,
+    span: Span,
+) -> Value {
+    if let Some(var) = block.signature.get_positional(0) {
+        if let Some(var_id) = &var.var_id {
+            stack.add_var(*var_id, val.clone());
+        }
+    }
+    match eval_block(engine_state, stack, block, val.into_pipeline_data()) {
+        Ok(pd) => pd.into_value(span),
+        Err(e) => Value::Error { error: e },
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_examples() {
+        use crate::test_examples;
+
+        test_examples(UpdateCells {})
+    }
+}