Make insert take in a block (#2265)

* Make insert take in a block

* Add some tests

crates/nu-cli/src/commands/insert.rs

@@ -1,8 +1,10 @@
+use crate::commands::classified::block::run_block;
 use crate::commands::WholeStreamCommand;
 use crate::context::CommandRegistry;
 use crate::prelude::*;
+use futures::stream::once;
 use nu_errors::ShellError;
-use nu_protocol::{ColumnPath, ReturnSuccess, Signature, SyntaxShape, UntaggedValue, Value};
+use nu_protocol::{ColumnPath, ReturnSuccess, Scope, Signature, SyntaxShape, UntaggedValue, Value};
 use nu_value_ext::ValueExt;
 
 pub struct Insert;
@@ -26,11 +28,7 @@ impl WholeStreamCommand for Insert {
                 SyntaxShape::ColumnPath,
                 "the column name to insert",
             )
-            .required(
-                "value",
-                SyntaxShape::String,
-                "the value to give the cell(s)",
-            )
+            .required("value", SyntaxShape::Any, "the value to give the cell(s)")
     }
 
     fn usage(&self) -> &str {
@@ -46,26 +44,111 @@ impl WholeStreamCommand for Insert {
     }
 }
 
-async fn insert(args: CommandArgs, registry: &CommandRegistry) -> Result<OutputStream, ShellError> {
-    let registry = registry.clone();
+async fn process_row(
+    scope: Arc<Scope>,
+    mut context: Arc<Context>,
+    input: Value,
+    mut value: Arc<Value>,
+    column: Arc<ColumnPath>,
+) -> Result<OutputStream, ShellError> {
+    let value = Arc::make_mut(&mut value);
 
-    let (InsertArgs { column, value }, input) = args.process(&registry).await?;
+    Ok(match value {
+        Value {
+            value: UntaggedValue::Block(block),
+            ..
+        } => {
+            let for_block = input.clone();
+            let input_stream = once(async { Ok(for_block) }).to_input_stream();
 
-    Ok(input
-        .map(move |row| match row {
+            let result = run_block(
+                &block,
+                Arc::make_mut(&mut context),
+                input_stream,
+                &input,
+                &scope.vars,
+                &scope.env,
+            )
+            .await;
+
+            match result {
+                Ok(mut stream) => {
+                    let errors = context.get_errors();
+                    if let Some(error) = errors.first() {
+                        return Err(error.clone());
+                    }
+
+                    match input {
+                        obj
+                        @
+                        Value {
+                            value: UntaggedValue::Row(_),
+                            ..
+                        } => {
+                            if let Some(result) = stream.next().await {
+                                match obj.insert_data_at_column_path(&column, result) {
+                                    Ok(v) => OutputStream::one(ReturnSuccess::value(v)),
+                                    Err(e) => OutputStream::one(Err(e)),
+                                }
+                            } else {
+                                OutputStream::empty()
+                            }
+                        }
+                        Value { tag, .. } => OutputStream::one(Err(ShellError::labeled_error(
+                            "Unrecognized type in stream",
+                            "original value",
+                            tag,
+                        ))),
+                    }
+                }
+                Err(e) => OutputStream::one(Err(e)),
+            }
+        }
+        _ => match input {
+            obj
+            @
             Value {
                 value: UntaggedValue::Row(_),
                 ..
-            } => Ok(ReturnSuccess::Value(
-                row.insert_data_at_column_path(&column, value.clone())?,
-            )),
-
-            Value { tag, .. } => Err(ShellError::labeled_error(
+            } => match obj.insert_data_at_column_path(&column, value.clone()) {
+                Ok(v) => OutputStream::one(ReturnSuccess::value(v)),
+                Err(e) => OutputStream::one(Err(e)),
+            },
+            Value { tag, .. } => OutputStream::one(Err(ShellError::labeled_error(
                 "Unrecognized type in stream",
                 "original value",
                 tag,
-            )),
+            ))),
+        },
+    })
+}
+
+async fn insert(
+    raw_args: CommandArgs,
+    registry: &CommandRegistry,
+) -> Result<OutputStream, ShellError> {
+    let registry = registry.clone();
+    let scope = Arc::new(raw_args.call_info.scope.clone());
+    let context = Arc::new(Context::from_raw(&raw_args, &registry));
+    let (InsertArgs { column, value }, input) = raw_args.process(&registry).await?;
+    let value = Arc::new(value);
+    let column = Arc::new(column);
+
+    Ok(input
+        .then(move |input| {
+            let scope = scope.clone();
+            let context = context.clone();
+            let value = value.clone();
+            let column = column.clone();
+
+            async {
+                match process_row(scope, context, input, value, column).await {
+                    Ok(s) => s,
+                    Err(e) => OutputStream::one(Err(e)),
+                }
+            }
         })
+        .flatten()
         .to_output_stream())
 }
 

crates/nu-cli/src/commands/update.rs

@@ -148,9 +148,9 @@ async fn update(
 
     Ok(input
         .then(move |input| {
-            let replacement = replacement.clone();
             let scope = scope.clone();
             let context = context.clone();
+            let replacement = replacement.clone();
             let field = field.clone();
 
             async {

crates/nu-cli/tests/commands/insert.rs

@@ -14,3 +14,27 @@ fn insert_plugin() {
 
     assert_eq!(actual.out, "1");
 }
+
+#[test]
+fn downcase_upcase() {
+    let actual = nu!(
+        cwd: ".", pipeline(
+        r#"
+        echo abcd | wrap downcase | insert upcase { echo $it.downcase | str upcase } | format "{downcase}{upcase}" 
+        "#
+    ));
+
+    assert_eq!(actual.out, "abcdABCD");
+}
+
+#[test]
+fn number_and_its_negative_equal_zero() {
+    let actual = nu!(
+        cwd: ".", pipeline(
+        r#"
+        echo 1..10 | wrap num | insert neg { = $it.num * -1 } | math sum | = $it.num + $it.neg
+        "#
+    ));
+
+    assert_eq!(actual.out, "0");
+}