diff --git a/crates/nu-command/src/filters/default.rs b/crates/nu-command/src/filters/default.rs index 51e6a0aadb..d424b80273 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::{ClosureEval, command_prelude::*}; 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.", @@ -43,8 +43,18 @@ impl Command for Default { call: &Call, input: PipelineData, ) -> Result { + let default_value: Value = call.req(engine_state, stack, 0)?; + let columns: Vec = call.rest(engine_state, stack, 1)?; let empty = call.has_flag(engine_state, stack, "empty")?; - default(engine_state, stack, call, input, empty) + default( + engine_state, + stack, + call, + input, + default_value, + empty, + columns, + ) } fn examples(&self) -> Vec { @@ -100,6 +110,25 @@ 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), + }), + ])), + }, ] } } @@ -109,43 +138,68 @@ fn default( stack: &mut Stack, call: &Call, input: PipelineData, + default_value: Value, default_when_empty: bool, + columns: Vec, ) -> Result { + let input_span = input.span().unwrap_or(call.head); + let mut default_value = DefaultValue::new(engine_state, stack, default_value); let metadata = input.metadata(); - let value: Value = call.req(engine_state, stack, 0)?; - let column: Option> = call.opt(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()); - } + // 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() { + if matches!(input, PipelineData::Value(Value::Record { .. }, _)) { + let record = input.into_value(input_span)?.into_record()?; + fill_record( + record, + input_span, + &mut default_value, + columns.as_slice(), + default_when_empty, + ) + .map(|x| x.into_pipeline_data_with_metadata(metadata)) + } else if matches!( + input, + PipelineData::ListStream(..) | PipelineData::Value(Value::List { .. }, _) + ) { + // Potential enhancement: add another branch for Value::List, + // and collect the iterator into a Result + // so we can preemptively return an error for collected lists + let head = call.head; + Ok(input + .into_iter() + .map(move |item| { + let span = item.span(); + if let Value::Record { val, .. } = item { + fill_record( + val.into_owned(), + span, + &mut default_value, + columns.as_slice(), + default_when_empty, + ) + .unwrap_or_else(|err| Value::error(err, head)) + } else { item } - _ => item, - }, - engine_state.signals(), - ) - .map(|x| x.set_metadata(metadata)) + }) + .into_pipeline_data_with_metadata(head, engine_state.signals().clone(), 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()) + default_value.pipeline_data() } else if default_when_empty && matches!(input, PipelineData::ListStream(..)) { let PipelineData::ListStream(ls, metadata) = input else { unreachable!() @@ -153,26 +207,86 @@ 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 default_value.pipeline_data(); } // 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) } } +/// A wrapper around the default value to handle closures and caching values +enum DefaultValue { + Uncalculated(Spanned), + Calculated(Value), +} + +impl DefaultValue { + fn new(engine_state: &EngineState, stack: &Stack, value: Value) -> Self { + let span = value.span(); + match value { + Value::Closure { val, .. } => { + let closure_eval = ClosureEval::new(engine_state, stack, *val); + DefaultValue::Uncalculated(closure_eval.into_spanned(span)) + } + _ => DefaultValue::Calculated(value), + } + } + + fn value(&mut self) -> Result { + match self { + DefaultValue::Uncalculated(closure) => { + let value = closure + .item + .run_with_input(PipelineData::Empty)? + .into_value(closure.span)?; + *self = DefaultValue::Calculated(value.clone()); + Ok(value) + } + DefaultValue::Calculated(value) => Ok(value.clone()), + } + } + + fn pipeline_data(&mut self) -> Result { + self.value().map(|x| x.into_pipeline_data()) + } +} + +/// Given a record, fill missing columns with a default value +fn fill_record( + mut record: Record, + span: Span, + default_value: &mut DefaultValue, + columns: &[String], + empty: bool, +) -> Result { + for col in columns { + if let Some(val) = record.get_mut(col) { + if matches!(val, Value::Nothing { .. }) || (empty && val.is_empty()) { + *val = default_value.value()?; + } + } else { + record.push(col.clone(), default_value.value()?); + } + } + Ok(Value::record(record, span)) +} + #[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]); } } diff --git a/crates/nu-command/tests/commands/default.rs b/crates/nu-command/tests/commands/default.rs index b17e334b12..75acf6d187 100644 --- a/crates/nu-command/tests/commands/default.rs +++ b/crates/nu-command/tests/commands/default.rs @@ -150,3 +150,97 @@ fn do_not_replace_non_empty_list_stream() { assert_eq!(actual.out, "2"); }) } + +#[test] +fn closure_eval_simple() { + let actual = nu!(r#"null | default { 1 }"#); + assert_eq!(actual.out, "1"); +} + +#[test] +fn closure_eval_complex() { + let actual = nu!(r#"null | default { seq 1 5 | math sum }"#); + assert_eq!(actual.out, "15"); +} + +#[test] +fn closure_eval_is_lazy() { + let actual = nu!(r#"1 | default { error make -u {msg: foo} }"#); + assert_eq!(actual.out, "1"); +} + +#[test] +fn column_closure_eval_is_lazy() { + let actual = nu!(r#"{a: 1} | default { error make -u {msg: foo} } a | get a"#); + assert_eq!(actual.out, "1"); +} + +#[test] +fn closure_eval_replace_empty_string() { + let actual = nu!(r#"'' | default --empty { 1 }"#); + assert_eq!(actual.out, "1"); +} + +#[test] +fn closure_eval_do_not_replace_empty_string() { + let actual = nu!(r#"'' | default { 1 }"#); + assert_eq!(actual.out, ""); +} + +#[test] +fn closure_eval_replace_empty_list() { + let actual = nu!(r#"[] | default --empty { 1 }"#); + assert_eq!(actual.out, "1"); +} + +#[test] +fn closure_eval_do_not_replace_empty_list() { + let actual = nu!(r#"[] | default { 1 } | length"#); + assert_eq!(actual.out, "0"); +} + +#[test] +fn closure_eval_replace_empty_record() { + let actual = nu!(r#"{} | default --empty { 1 }"#); + assert_eq!(actual.out, "1"); +} + +#[test] +fn closure_eval_do_not_replace_empty_record() { + let actual = nu!(r#"{} | default { 1 } | columns | length"#); + assert_eq!(actual.out, "0"); +} + +#[test] +fn closure_eval_add_missing_column_record() { + let actual = nu!(r#" + {a: 1} | default { 2 } b | get b + "#); + assert_eq!(actual.out, "2"); +} + +#[test] +fn closure_eval_add_missing_column_table() { + let actual = nu!(r#" + [{a: 1, b: 2}, {b: 4}] | default { 3 } a | get a | to json -r + "#); + assert_eq!(actual.out, "[1,3]"); +} + +#[test] +fn closure_eval_replace_empty_column() { + let actual = nu!(r#"{a: ''} | default -e { 1 } a | get a"#); + assert_eq!(actual.out, "1"); +} + +#[test] +fn replace_multiple_columns() { + let actual = nu!(r#"{a: ''} | default -e 1 a b | values | to json -r"#); + assert_eq!(actual.out, "[1,1]"); +} + +#[test] +fn return_closure_value() { + let actual = nu!(r#"null | default { {||} }"#); + assert!(actual.out.starts_with("closure")); +}