diff --git a/crates/nu-command/src/generators/generate.rs b/crates/nu-command/src/generators/generate.rs index 8d4f48d3fb..2b2cf838f4 100644 --- a/crates/nu-command/src/generators/generate.rs +++ b/crates/nu-command/src/generators/generate.rs @@ -12,12 +12,12 @@ impl Command for Generate { fn signature(&self) -> Signature { Signature::build("generate") .input_output_types(vec![(Type::Nothing, Type::List(Box::new(Type::Any)))]) - .required("initial", SyntaxShape::Any, "Initial value.") .required( "closure", SyntaxShape::Closure(Some(vec![SyntaxShape::Any])), "Generator function.", ) + .optional("initial", SyntaxShape::Any, "Initial value.") .allow_variants_without_examples(true) .category(Category::Generators) } @@ -41,7 +41,7 @@ used as the next argument to the closure, otherwise generation stops. fn examples(&self) -> Vec { vec![ Example { - example: "generate 0 {|i| if $i <= 10 { {out: $i, next: ($i + 2)} }}", + example: "generate {|i| if $i <= 10 { {out: $i, next: ($i + 2)} }} 0", description: "Generate a sequence of numbers", result: Some(Value::list( vec![ @@ -57,10 +57,17 @@ used as the next argument to the closure, otherwise generation stops. }, Example { example: - "generate [0, 1] {|fib| {out: $fib.0, next: [$fib.1, ($fib.0 + $fib.1)]} }", + "generate {|fib| {out: $fib.0, next: [$fib.1, ($fib.0 + $fib.1)]} } [0, 1]", description: "Generate a continuous stream of Fibonacci numbers", result: None, }, + Example { + example: + "generate {|fib=[0, 1]| {out: $fib.0, next: [$fib.1, ($fib.0 + $fib.1)]} }", + description: + "Generate a continuous stream of Fibonacci numbers, using default parameters", + result: None, + }, ] } @@ -72,15 +79,15 @@ used as the next argument to the closure, otherwise generation stops. _input: PipelineData, ) -> Result { let head = call.head; - let initial: Value = call.req(engine_state, stack, 0)?; - let closure: Closure = call.req(engine_state, stack, 1)?; - + let closure: Closure = call.req(engine_state, stack, 0)?; + let initial: Option = call.opt(engine_state, stack, 1)?; + let block = engine_state.get_block(closure.block_id); let mut closure = ClosureEval::new(engine_state, stack, closure); // A type of Option is used to represent state. Invocation // will stop on None. Using Option allows functions to output // one final value before stopping. - let mut state = Some(initial); + let mut state = Some(get_initial_state(initial, &block.signature, call.head)?); let iter = std::iter::from_fn(move || { let arg = state.take()?; @@ -170,6 +177,38 @@ used as the next argument to the closure, otherwise generation stops. } } +fn get_initial_state( + initial: Option, + signature: &Signature, + span: Span, +) -> Result { + match initial { + Some(v) => Ok(v), + None => { + // the initial state should be referred from signature + if !signature.optional_positional.is_empty() + && signature.optional_positional[0].default_value.is_some() + { + Ok(signature.optional_positional[0] + .default_value + .clone() + .expect("Already checked default value")) + } else { + Err(ShellError::GenericError { + error: "The initial value is missing".to_string(), + msg: "Missing initial value".to_string(), + span: Some(span), + help: Some( + "Provide an value as an argument to generate, or assign a default value to the closure parameter" + .to_string(), + ), + inner: vec![], + }) + } + } + } +} + #[cfg(test)] mod test { use super::*; diff --git a/crates/nu-command/tests/commands/generate.rs b/crates/nu-command/tests/commands/generate.rs index 9dc2142515..9ae3b8823e 100644 --- a/crates/nu-command/tests/commands/generate.rs +++ b/crates/nu-command/tests/commands/generate.rs @@ -3,7 +3,7 @@ use nu_test_support::{nu, pipeline}; #[test] fn generate_no_next_break() { let actual = nu!( - "generate 1 {|x| if $x == 3 { {out: $x}} else { {out: $x, next: ($x + 1)} }} | to nuon" + "generate {|x| if $x == 3 { {out: $x}} else { {out: $x, next: ($x + 1)} }} 1 | to nuon" ); assert_eq!(actual.out, "[1, 2, 3]"); @@ -11,7 +11,7 @@ fn generate_no_next_break() { #[test] fn generate_null_break() { - let actual = nu!("generate 1 {|x| if $x <= 3 { {out: $x, next: ($x + 1)} }} | to nuon"); + let actual = nu!("generate {|x| if $x <= 3 { {out: $x, next: ($x + 1)} }} 1 | to nuon"); assert_eq!(actual.out, "[1, 2, 3]"); } @@ -20,13 +20,13 @@ fn generate_null_break() { fn generate_allows_empty_output() { let actual = nu!(pipeline( r#" - generate 0 {|x| + generate {|x| if $x == 1 { {next: ($x + 1)} } else if $x < 3 { {out: $x, next: ($x + 1)} } - } | to nuon + } 0 | to nuon "# )); @@ -37,11 +37,11 @@ fn generate_allows_empty_output() { fn generate_allows_no_output() { let actual = nu!(pipeline( r#" - generate 0 {|x| + generate {|x| if $x < 3 { {next: ($x + 1)} } - } | to nuon + } 0 | to nuon "# )); @@ -52,7 +52,7 @@ fn generate_allows_no_output() { fn generate_allows_null_state() { let actual = nu!(pipeline( r#" - generate 0 {|x| + generate {|x| if $x == null { {out: "done"} } else if $x < 1 { @@ -60,7 +60,7 @@ fn generate_allows_null_state() { } else { {out: "stopping", next: null} } - } | to nuon + } 0 | to nuon "# )); @@ -71,7 +71,42 @@ fn generate_allows_null_state() { fn generate_allows_null_output() { let actual = nu!(pipeline( r#" - generate 0 {|x| + generate {|x| + if $x == 3 { + {out: "done"} + } else { + {out: null, next: ($x + 1)} + } + } 0 | to nuon + "# + )); + + assert_eq!(actual.out, "[null, null, null, done]"); +} + +#[test] +fn generate_disallows_extra_keys() { + let actual = nu!("generate {|x| {foo: bar, out: $x}} 0 "); + assert!(actual.err.contains("Invalid block return")); +} + +#[test] +fn generate_disallows_list() { + let actual = nu!("generate {|x| [$x, ($x + 1)]} 0 "); + assert!(actual.err.contains("Invalid block return")); +} + +#[test] +fn generate_disallows_primitive() { + let actual = nu!("generate {|x| 1} 0"); + assert!(actual.err.contains("Invalid block return")); +} + +#[test] +fn generate_allow_default_parameter() { + let actual = nu!(pipeline( + r#" + generate {|x = 0| if $x == 3 { {out: "done"} } else { @@ -82,22 +117,34 @@ fn generate_allows_null_output() { )); assert_eq!(actual.out, "[null, null, null, done]"); + + // if initial is given, use initial value + let actual = nu!(pipeline( + r#" + generate {|x = 0| + if $x == 3 { + {out: "done"} + } else { + {out: null, next: ($x + 1)} + } + } 1 | to nuon + "# + )); + assert_eq!(actual.out, "[null, null, done]"); } #[test] -fn generate_disallows_extra_keys() { - let actual = nu!("generate 0 {|x| {foo: bar, out: $x}}"); - assert!(actual.err.contains("Invalid block return")); -} - -#[test] -fn generate_disallows_list() { - let actual = nu!("generate 0 {|x| [$x, ($x + 1)]}"); - assert!(actual.err.contains("Invalid block return")); -} - -#[test] -fn generate_disallows_primitive() { - let actual = nu!("generate 0 {|x| 1}"); - assert!(actual.err.contains("Invalid block return")); +fn generate_raise_error_on_no_default_parameter_closure_and_init_val() { + let actual = nu!(pipeline( + r#" + generate {|x| + if $x == 3 { + {out: "done"} + } else { + {out: null, next: ($x + 1)} + } + } | to nuon + "# + )); + assert!(actual.err.contains("The initial value is missing")); }