diff --git a/crates/nu-cli/src/commands/insert.rs b/crates/nu-cli/src/commands/insert.rs index 3039a8c15a..e938d54ac5 100644 --- a/crates/nu-cli/src/commands/insert.rs +++ b/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 { - let registry = registry.clone(); +async fn process_row( + scope: Arc, + mut context: Arc, + input: Value, + mut value: Arc, + column: Arc, +) -> Result { + let value = Arc::make_mut(&mut value); - let (InsertArgs { column, value }, input) = args.process(®istry).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 { + let registry = registry.clone(); + let scope = Arc::new(raw_args.call_info.scope.clone()); + let context = Arc::new(Context::from_raw(&raw_args, ®istry)); + let (InsertArgs { column, value }, input) = raw_args.process(®istry).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()) } diff --git a/crates/nu-cli/src/commands/update.rs b/crates/nu-cli/src/commands/update.rs index 2590d42bb4..a74d256942 100644 --- a/crates/nu-cli/src/commands/update.rs +++ b/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 { diff --git a/crates/nu-cli/tests/commands/insert.rs b/crates/nu-cli/tests/commands/insert.rs index 9323256cd8..ecd8be7b21 100644 --- a/crates/nu-cli/tests/commands/insert.rs +++ b/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"); +}