diff --git a/crates/nu-command/src/filters/default.rs b/crates/nu-command/src/filters/default.rs index f0dc641d5f..36d4f0e5b6 100644 --- a/crates/nu-command/src/filters/default.rs +++ b/crates/nu-command/src/filters/default.rs @@ -1,4 +1,4 @@ -use nu_engine::command_prelude::*; +use nu_engine::{command_prelude::*, ClosureEval, ClosureEvalOnce}; use nu_protocol::{ListStream, Signals}; #[derive(Clone)] @@ -19,7 +19,7 @@ impl Command for Default { SyntaxShape::Any, "The value to use as a default.", ) - .optional( + .rest( "column name", SyntaxShape::String, "The name of the column.", @@ -101,10 +101,92 @@ impl Command for Default { }), ])), }, + Example { + description: r#"Generate a default value from a closure"#, + example: "null | default { 1 + 2 }", + result: Some(Value::test_int(3)), + }, + Example { + description: r#"Fill missing column values based on other columns"#, + example: r#"[{a:1 b:2} {b:1}] | upsert a {|rc| default { $rc.b + 1 } }"#, + result: Some(Value::test_list(vec![ + Value::test_record(record! { + "a" => Value::test_int(1), + "b" => Value::test_int(2), + }), + Value::test_record(record! { + "a" => Value::test_int(2), + "b" => Value::test_int(1), + }), + ])), + }, ] } } +fn eval_default( + engine_state: &EngineState, + stack: &mut Stack, + default_value: Value, +) -> Result { + match &default_value { + Value::Closure { val, .. } => { + let closure = ClosureEvalOnce::new(engine_state, stack, *val.clone()); + closure.run_with_input(PipelineData::Empty) + } + _ => Ok(default_value.into_pipeline_data()), + } +} + +fn default_record_columns( + record: &mut Record, + default_value: Spanned, + columns: &[String], + empty: bool, + engine_state: &EngineState, + stack: &mut Stack, + calculated_value: &mut Option, +) -> Result { + if let Value::Closure { val: closure, .. } = &default_value.item { + // Cache the value of the closure to avoid running it multiple times + let mut closure = ClosureEval::new(engine_state, stack, *closure.clone()); + for col in columns { + if let Some(val) = record.get_mut(col) { + if matches!(val, Value::Nothing { .. }) || (empty && val.is_empty()) { + if let Some(ref new_value) = calculated_value { + *val = new_value.clone(); + } else { + let new_value = closure + .run_with_input(PipelineData::Empty)? + .into_value(default_value.span)?; + *calculated_value = Some(new_value.clone()); + *val = new_value; + } + } + } else if let Some(ref new_value) = calculated_value { + record.push(col.clone(), new_value.clone()); + } else { + let new_value = closure + .run_with_input(PipelineData::Empty)? + .into_value(default_value.span)?; + *calculated_value = Some(new_value.clone()); + record.push(col.clone(), new_value); + } + } + } else { + for col in columns { + if let Some(val) = record.get_mut(col) { + if matches!(val, Value::Nothing { .. }) || (empty && val.is_empty()) { + *val = default_value.item.clone(); + } + } else { + record.push(col.clone(), default_value.item.clone()); + } + } + } + Ok(Value::record(record.clone(), Span::unknown()).into_pipeline_data()) +} + fn default( engine_state: &EngineState, stack: &mut Stack, @@ -112,41 +194,81 @@ fn default( input: PipelineData, default_when_empty: bool, ) -> Result { + let input_span = input.span().unwrap_or_else(Span::unknown); let metadata = input.metadata(); - let value: Value = call.req(engine_state, stack, 0)?; - let column: Option> = call.opt(engine_state, stack, 1)?; + let default_value: Spanned = call.req(engine_state, stack, 0)?; + let columns: Vec = call.rest(engine_state, stack, 1)?; - if let Some(column) = column { - input - .map( - move |mut item| match item { - Value::Record { - val: ref mut record, - .. - } => { - let record = record.to_mut(); - if let Some(val) = record.get_mut(&column.item) { - if matches!(val, Value::Nothing { .. }) - || (default_when_empty && val.is_empty()) - { - *val = value.clone(); - } - } else { - record.push(column.item.clone(), value.clone()); - } - - item - } - _ => item, - }, - engine_state.signals(), + // If user supplies columns, check if input is a record or list of records + // and set the default value for the specified record columns + if !columns.is_empty() { + // Single record arm + if matches!(input, PipelineData::Value(Value::Record { .. }, _)) { + let Value::Record { + val: ref mut record, + .. + } = input.into_value(input_span)? + else { + unreachable!() + }; + let record = record.to_mut(); + default_record_columns( + record, + default_value, + columns.as_slice(), + default_when_empty, + engine_state, + stack, + &mut None, ) .map(|x| x.set_metadata(metadata)) + // ListStream arm + } else if matches!(input, PipelineData::ListStream(..)) + || matches!(input, PipelineData::Value(Value::List { .. }, _)) + { + let mut calculated_value: Option = None; + let mut output_list: Vec = vec![]; + for mut item in input { + if let Value::Record { + val: ref mut record, + internal_span, + } = item + { + let item = default_record_columns( + record.to_mut(), + default_value.clone(), + columns.as_slice(), + default_when_empty, + engine_state, + stack, + &mut calculated_value, + )?; + output_list.push(item.into_value(internal_span)?); + } else { + output_list.push(item); + } + } + let ls = ListStream::new( + output_list.into_iter(), + call.head, + engine_state.signals().clone(), + ); + Ok(PipelineData::ListStream(ls, metadata)) + // If columns are given, but input does not use columns, return an error + } else { + Err(ShellError::PipelineMismatch { + exp_input_type: "record, table".to_string(), + dst_span: input_span, + src_span: input_span, + }) + } + // Otherwise, if no column name is given, check if value is null + // or an empty string, list, or record when --empty is passed } else if input.is_nothing() || (default_when_empty && matches!(input, PipelineData::Value(ref value, _) if value.is_empty())) { - Ok(value.into_pipeline_data()) + eval_default(engine_state, stack, default_value.item) } else if default_when_empty && matches!(input, PipelineData::ListStream(..)) { let PipelineData::ListStream(ls, metadata) = input else { unreachable!() @@ -154,13 +276,14 @@ fn default( let span = ls.span(); let mut stream = ls.into_inner().peekable(); if stream.peek().is_none() { - return Ok(value.into_pipeline_data()); + return eval_default(engine_state, stack, default_value.item); } // stream's internal state already preserves the original signals config, so if this // Signals::empty list stream gets interrupted it will be caught by the underlying iterator let ls = ListStream::new(stream, span, Signals::empty()); Ok(PipelineData::ListStream(ls, metadata)) + // Otherwise, return the input as is } else { Ok(input) } @@ -168,12 +291,14 @@ fn default( #[cfg(test)] mod test { + use crate::Upsert; + use super::*; #[test] fn test_examples() { - use crate::test_examples; + use crate::test_examples_with_commands; - test_examples(Default {}) + test_examples_with_commands(Default {}, &[&Upsert]); } }