use crate::command_registry::CommandRegistry; use crate::commands::classified::block::run_block; use crate::commands::WholeStreamCommand; use crate::prelude::*; use nu_errors::ShellError; use nu_protocol::{ ColumnPath, Primitive, ReturnSuccess, Scope, Signature, SyntaxShape, UntaggedValue, Value, }; use nu_value_ext::ValueExt; use futures::stream::once; use indexmap::indexmap; pub struct Command; #[derive(Deserialize)] pub struct Arguments { column: ColumnPath, value: Value, } #[async_trait] impl WholeStreamCommand for Command { fn name(&self) -> &str { "insert" } fn signature(&self) -> Signature { Signature::build("insert") .required( "column", SyntaxShape::ColumnPath, "the column name to insert", ) .required("value", SyntaxShape::Any, "the value to give the cell(s)") } fn usage(&self) -> &str { "Insert a new column with a given value." } async fn run( &self, args: CommandArgs, registry: &CommandRegistry, ) -> Result { insert(args, registry).await } fn examples(&self) -> Vec { vec![Example { description: "Insert a column with a value", example: "echo [[author, commits]; ['Andrés', 1]] | insert branches 5", result: Some(vec![UntaggedValue::row(indexmap! { "author".to_string() => Value::from("Andrés"), "commits".to_string() => UntaggedValue::int(1).into(), "branches".to_string() => UntaggedValue::int(5).into(), }) .into()]), },Example { description: "Use in block form for more involved insertion logic", example: "echo [[author, lucky_number]; ['Yehuda', 4]] | insert success { = $it.lucky_number * 10 }", result: Some(vec![UntaggedValue::row(indexmap! { "author".to_string() => Value::from("Yehuda"), "lucky_number".to_string() => UntaggedValue::int(4).into(), "success".to_string() => UntaggedValue::int(40).into(), }) .into()]), }] } } async fn process_row( scope: Arc, mut context: Arc, input: Value, mut value: Arc, field: Arc, ) -> Result { let value = Arc::make_mut(&mut value); Ok(match value { Value { value: UntaggedValue::Block(block), tag: block_tag, } => { let for_block = input.clone(); let input_stream = once(async { Ok(for_block) }).to_input_stream(); let scope = Scope::append_it(scope, input.clone()); let result = run_block(&block, Arc::make_mut(&mut context), input_stream, scope).await; match result { Ok(mut stream) => { let values = stream.drain_vec().await; let errors = context.get_errors(); if let Some(error) = errors.first() { return Err(error.clone()); } let result = if values.len() == 1 { let value = values .get(0) .ok_or_else(|| ShellError::unexpected("No value to insert with."))?; Value { value: value.value.clone(), tag: input.tag.clone(), } } else if values.is_empty() { UntaggedValue::nothing().into_value(&input.tag) } else { UntaggedValue::table(&values).into_value(&input.tag) }; match input { obj @ Value { value: UntaggedValue::Row(_), .. } => match obj.insert_data_at_column_path(&field, result) { Ok(v) => OutputStream::one(ReturnSuccess::value(v)), Err(e) => OutputStream::one(Err(e)), }, _ => OutputStream::one(Err(ShellError::labeled_error( "Unrecognized type in stream", "original value", block_tag.clone(), ))), } } Err(e) => OutputStream::one(Err(e)), } } value => match input { Value { value: UntaggedValue::Primitive(Primitive::Nothing), .. } => match scope .it() .unwrap_or_else(|| UntaggedValue::nothing().into_untagged_value()) .insert_data_at_column_path(&field, value.clone()) { Ok(v) => OutputStream::one(ReturnSuccess::value(v)), Err(e) => OutputStream::one(Err(e)), }, _ => match input.insert_data_at_column_path(&field, value.clone()) { Ok(v) => OutputStream::one(ReturnSuccess::value(v)), Err(e) => OutputStream::one(Err(e)), }, }, }) } async fn insert( raw_args: CommandArgs, registry: &CommandRegistry, ) -> Result { let registry = registry.clone(); let scope = raw_args.call_info.scope.clone(); let context = Arc::new(EvaluationContext::from_raw(&raw_args, ®istry)); let (Arguments { 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()) }