From e76451866dd6b70755d7f318a329c6631edb974a Mon Sep 17 00:00:00 2001
From: JT <547158+jntrnr@users.noreply.github.com>
Date: Sun, 14 Nov 2021 12:02:54 +1300
Subject: [PATCH] 'update' command (#333)

---
 crates/nu-command/src/default_context.rs |   1 +
 crates/nu-command/src/filters/mod.rs     |   2 +
 crates/nu-command/src/filters/update.rs  | 120 +++++++++++++++++++++++
 3 files changed, 123 insertions(+)
 create mode 100644 crates/nu-command/src/filters/update.rs

diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs
index e85e93eeee..3db8996467 100644
--- a/crates/nu-command/src/default_context.rs
+++ b/crates/nu-command/src/default_context.rs
@@ -119,6 +119,7 @@ pub fn create_default_context() -> EngineState {
             ToJson,
             Touch,
             Use,
+            Update,
             Where,
             WithEnv,
             Wrap,
diff --git a/crates/nu-command/src/filters/mod.rs b/crates/nu-command/src/filters/mod.rs
index ea9495a1fc..bf93145d5d 100644
--- a/crates/nu-command/src/filters/mod.rs
+++ b/crates/nu-command/src/filters/mod.rs
@@ -10,6 +10,7 @@ mod range;
 mod reverse;
 mod select;
 mod shuffle;
+mod update;
 mod where_;
 mod wrap;
 mod zip;
@@ -26,6 +27,7 @@ pub use range::Range;
 pub use reverse::Reverse;
 pub use select::Select;
 pub use shuffle::Shuffle;
+pub use update::Update;
 pub use where_::Where;
 pub use wrap::Wrap;
 pub use zip::Zip;
diff --git a/crates/nu-command/src/filters/update.rs b/crates/nu-command/src/filters/update.rs
new file mode 100644
index 0000000000..5e5b120ae2
--- /dev/null
+++ b/crates/nu-command/src/filters/update.rs
@@ -0,0 +1,120 @@
+use nu_engine::{eval_block, CallExt};
+use nu_protocol::ast::{Call, CellPath};
+use nu_protocol::engine::{Command, EngineState, Stack};
+use nu_protocol::{
+    Example, IntoPipelineData, PipelineData, ShellError, Signature, Span, SyntaxShape, Value,
+};
+
+#[derive(Clone)]
+pub struct Update;
+
+impl Command for Update {
+    fn name(&self) -> &str {
+        "update"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build("update")
+            .required(
+                "field",
+                SyntaxShape::CellPath,
+                "the name of the column to update",
+            )
+            .required(
+                "replacement value",
+                SyntaxShape::Any,
+                "the new value to give the cell(s)",
+            )
+    }
+
+    fn usage(&self) -> &str {
+        "Update an existing column to have a new value."
+    }
+
+    fn run(
+        &self,
+        engine_state: &EngineState,
+        stack: &mut Stack,
+        call: &Call,
+        input: PipelineData,
+    ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
+        update(engine_state, stack, call, input)
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![Example {
+            description: "Update a column value",
+            example: "echo {'name': 'nu', 'stars': 5} | update name 'Nushell'",
+            result: Some(Value::Record { cols: vec!["name".into(), "stars".into()], vals: vec![Value::test_string("Nushell"), Value::test_int(5)], span: Span::unknown()}),
+        }, Example {
+            description: "Use in block form for more involved updating logic",
+            example: "echo [[project, authors]; ['nu', ['Andrés', 'JT', 'Yehuda']]] | update authors { get authors | str collect ',' }",
+            result: Some(Value::List { vals: vec![Value::Record { cols: vec!["project".into(), "authors".into()], vals: vec![Value::test_string("nu"), Value::test_string("Andrés,JT,Yehuda")], span: Span::unknown()}], span: Span::unknown()}),
+        }]
+    }
+}
+
+fn update(
+    engine_state: &EngineState,
+    stack: &mut Stack,
+    call: &Call,
+    input: PipelineData,
+) -> Result<PipelineData, ShellError> {
+    let span = call.head;
+
+    let cell_path: CellPath = call.req(engine_state, stack, 0)?;
+    let replacement: Value = call.req(engine_state, stack, 1)?;
+    let engine_state = engine_state.clone();
+    let ctrlc = engine_state.ctrlc.clone();
+
+    // Replace is a block, so set it up and run it instead of using it as the replacement
+    if let Ok(block_id) = replacement.as_block() {
+        let block = engine_state.get_block(block_id).clone();
+
+        let mut stack = stack.collect_captures(&block.captures);
+
+        input.map(
+            move |mut input| {
+                if let Some(var) = block.signature.get_positional(0) {
+                    if let Some(var_id) = &var.var_id {
+                        stack.add_var(*var_id, input.clone())
+                    }
+                }
+
+                let output = eval_block(
+                    &engine_state,
+                    &mut stack,
+                    &block,
+                    input.clone().into_pipeline_data(),
+                );
+
+                match output {
+                    Ok(pd) => {
+                        if let Err(e) =
+                            input.replace_data_at_cell_path(&cell_path.members, pd.into_value(span))
+                        {
+                            return Value::Error { error: e };
+                        }
+
+                        input
+                    }
+                    Err(e) => Value::Error { error: e },
+                }
+            },
+            ctrlc,
+        )
+    } else {
+        input.map(
+            move |mut input| {
+                let replacement = replacement.clone();
+
+                if let Err(e) = input.replace_data_at_cell_path(&cell_path.members, replacement) {
+                    return Value::Error { error: e };
+                }
+
+                input
+            },
+            ctrlc,
+        )
+    }
+}