diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs index 7b11a62c69..04ac32873b 100644 --- a/crates/nu-command/src/default_context.rs +++ b/crates/nu-command/src/default_context.rs @@ -84,6 +84,7 @@ pub fn create_default_context(cwd: impl AsRef) -> EngineState { DropColumn, DropNth, Each, + EachWhile, Empty, Every, Find, diff --git a/crates/nu-command/src/filters/each_while.rs b/crates/nu-command/src/filters/each_while.rs new file mode 100644 index 0000000000..18f4469642 --- /dev/null +++ b/crates/nu-command/src/filters/each_while.rs @@ -0,0 +1,233 @@ +use nu_engine::{eval_block, CallExt}; +use nu_protocol::ast::Call; +use nu_protocol::engine::{CaptureBlock, Command, EngineState, Stack}; +use nu_protocol::{ + Category, Example, IntoInterruptiblePipelineData, IntoPipelineData, PipelineData, Signature, + Span, SyntaxShape, Value, +}; + +#[derive(Clone)] +pub struct EachWhile; + +impl Command for EachWhile { + fn name(&self) -> &str { + "each while" + } + + fn usage(&self) -> &str { + "Run a block on each element of input until a $nothing is found" + } + + fn search_terms(&self) -> Vec<&str> { + vec!["for", "loop", "iterate", "while"] + } + + fn signature(&self) -> nu_protocol::Signature { + Signature::build(self.name()) + .required( + "block", + SyntaxShape::Block(Some(vec![SyntaxShape::Any])), + "the block to run", + ) + .switch("numbered", "iterate with an index", Some('n')) + .category(Category::Filters) + } + + fn examples(&self) -> Vec { + let stream_test_1 = vec![ + Value::Int { + val: 1, + span: Span::test_data(), + }, + Value::Int { + val: 2, + span: Span::test_data(), + }, + ]; + + vec![ + Example { + example: "[1 2 3] | each while { |it| if $it < 3 {$it} else {$nothing} }", + description: "Multiplies elements in list", + result: Some(Value::List { + vals: stream_test_1, + span: Span::test_data(), + }), + }, + Example { + example: r#"[1 2 3] | each while -n { |it| if $it.item < 2 { $"value ($it.item) at ($it.index)!"} else { $nothing } }"#, + description: "Iterate over each element, print the matching value and its index", + result: Some(Value::List { + vals: vec![Value::String { + val: "value 1 at 0!".to_string(), + span: Span::test_data(), + }], + span: Span::test_data(), + }), + }, + ] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + let capture_block: CaptureBlock = call.req(engine_state, stack, 0)?; + let numbered = call.has_flag("numbered"); + + let metadata = input.metadata(); + let ctrlc = engine_state.ctrlc.clone(); + let engine_state = engine_state.clone(); + let block = engine_state.get_block(capture_block.block_id).clone(); + let mut stack = stack.captures_to_stack(&capture_block.captures); + let orig_env_vars = stack.env_vars.clone(); + let orig_env_hidden = stack.env_hidden.clone(); + let span = call.head; + let redirect_stdout = call.redirect_stdout; + let redirect_stderr = call.redirect_stderr; + + match input { + PipelineData::Value(Value::Range { .. }, ..) + | PipelineData::Value(Value::List { .. }, ..) + | PipelineData::ListStream { .. } => Ok(input + .into_iter() + .enumerate() + .map_while(move |(idx, x)| { + stack.with_env(&orig_env_vars, &orig_env_hidden); + + if let Some(var) = block.signature.get_positional(0) { + if let Some(var_id) = &var.var_id { + if numbered { + stack.add_var( + *var_id, + Value::Record { + cols: vec!["index".into(), "item".into()], + vals: vec![ + Value::Int { + val: idx as i64, + span, + }, + x.clone(), + ], + span, + }, + ); + } else { + stack.add_var(*var_id, x.clone()); + } + } + } + + match eval_block( + &engine_state, + &mut stack, + &block, + x.into_pipeline_data(), + redirect_stdout, + redirect_stderr, + ) { + Ok(v) => { + let value = v.into_value(span); + if value.is_nothing() { + None + } else { + Some(value) + } + } + Err(_) => None, + } + }) + .into_pipeline_data(ctrlc)), + PipelineData::ExternalStream { stdout: None, .. } => Ok(PipelineData::new(call.head)), + PipelineData::ExternalStream { + stdout: Some(stream), + .. + } => Ok(stream + .into_iter() + .enumerate() + .map_while(move |(idx, x)| { + stack.with_env(&orig_env_vars, &orig_env_hidden); + + let x = match x { + Ok(x) => x, + Err(_) => return None, + }; + + if let Some(var) = block.signature.get_positional(0) { + if let Some(var_id) = &var.var_id { + if numbered { + stack.add_var( + *var_id, + Value::Record { + cols: vec!["index".into(), "item".into()], + vals: vec![ + Value::Int { + val: idx as i64, + span, + }, + x.clone(), + ], + span, + }, + ); + } else { + stack.add_var(*var_id, x.clone()); + } + } + } + + match eval_block( + &engine_state, + &mut stack, + &block, + x.into_pipeline_data(), + redirect_stdout, + redirect_stderr, + ) { + Ok(v) => { + let value = v.into_value(span); + if value.is_nothing() { + None + } else { + Some(value) + } + } + Err(_) => None, + } + }) + .into_pipeline_data(ctrlc)), + PipelineData::Value(x, ..) => { + if let Some(var) = block.signature.get_positional(0) { + if let Some(var_id) = &var.var_id { + stack.add_var(*var_id, x.clone()); + } + } + + eval_block( + &engine_state, + &mut stack, + &block, + x.into_pipeline_data(), + redirect_stdout, + redirect_stderr, + ) + } + } + .map(|x| x.set_metadata(metadata)) + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_examples() { + use crate::test_examples; + + test_examples(EachWhile {}) + } +} diff --git a/crates/nu-command/src/filters/mod.rs b/crates/nu-command/src/filters/mod.rs index c92ee16e36..8ce7788a13 100644 --- a/crates/nu-command/src/filters/mod.rs +++ b/crates/nu-command/src/filters/mod.rs @@ -7,6 +7,7 @@ mod compact; mod default; mod drop; mod each; +mod each_while; mod empty; mod every; mod find; @@ -57,6 +58,7 @@ pub use compact::Compact; pub use default::Default; pub use drop::*; pub use each::Each; +pub use each_while::EachWhile; pub use empty::Empty; pub use every::Every; pub use find::Find; diff --git a/src/tests/test_engine.rs b/src/tests/test_engine.rs index d0553e0cd4..17ed63c3e1 100644 --- a/src/tests/test_engine.rs +++ b/src/tests/test_engine.rs @@ -53,7 +53,7 @@ fn in_and_if_else() -> TestResult { #[test] fn help_works_with_missing_requirements() -> TestResult { - run_test(r#"each --help | lines | length"#, "29") + run_test(r#"each --help | lines | length"#, "32") } #[test]