Merge 9d510ff66ce4d6f41a2684ff8b7f73a5e7d89f7f into a0d7c1a4fde4a6b052f536375f3790676fd7bc31

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<PipelineData, ShellError> {
+    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<Value>,
+    columns: &[String],
+    empty: bool,
+    engine_state: &EngineState,
+    stack: &mut Stack,
+    calculated_value: &mut Option<Value>,
+) -> Result<PipelineData, ShellError> {
+    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<PipelineData, ShellError> {
+    let input_span = input.span().unwrap_or_else(Span::unknown);
     let metadata = input.metadata();
-    let value: Value = call.req(engine_state, stack, 0)?;
+    let default_value: Spanned<Value> = call.req(engine_state, stack, 0)?;
-    let column: Option<Spanned<String>> = call.opt(engine_state, stack, 1)?;
+    let columns: Vec<String> = call.rest(engine_state, stack, 1)?;
 
-    if let Some(column) = column {
+    // If user supplies columns, check if input is a record or list of records
-        input
+    // and set the default value for the specified record columns
-            .map(
+    if !columns.is_empty() {
-                move |mut item| match item {
+        // Single record arm
-                    Value::Record {
+        if matches!(input, PipelineData::Value(Value::Record { .. }, _)) {
-                        val: ref mut record,
+            let Value::Record {
-                        ..
+                val: ref mut record,
-                    } => {
+                ..
-                        let record = record.to_mut();
+            } = input.into_value(input_span)?
-                        if let Some(val) = record.get_mut(&column.item) {
+            else {
-                            if matches!(val, Value::Nothing { .. })
+                unreachable!()
-                                || (default_when_empty && val.is_empty())
+            };
-                            {
+            let record = record.to_mut();
-                                *val = value.clone();
+            default_record_columns(
-                            }
+                record,
-                        } else {
+                default_value,
-                            record.push(column.item.clone(), value.clone());
+                columns.as_slice(),
-                        }
+                default_when_empty,
-
+                engine_state,
-                        item
+                stack,
-                    }
+                &mut None,
-                    _ => item,
-                },
-                engine_state.signals(),
             )
             .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<Value> = None;
+            let mut output_list: Vec<Value> = 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]);
     }
 }