diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs index 05b198719f..f8f4710cb2 100644 --- a/crates/nu-command/src/default_context.rs +++ b/crates/nu-command/src/default_context.rs @@ -66,6 +66,7 @@ pub fn create_default_context(cwd: impl AsRef) -> EngineState { DropColumn, DropNth, Each, + EachGroup, Empty, Every, Find, diff --git a/crates/nu-command/src/filters/each_group.rs b/crates/nu-command/src/filters/each_group.rs new file mode 100644 index 0000000000..0780888158 --- /dev/null +++ b/crates/nu-command/src/filters/each_group.rs @@ -0,0 +1,161 @@ +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, Spanned, SyntaxShape, Value, +}; + +#[derive(Clone)] +pub struct EachGroup; + +impl Command for EachGroup { + fn name(&self) -> &str { + "each group" + } + + fn signature(&self) -> Signature { + Signature::build("each group") + .required("group_size", SyntaxShape::Int, "the size of each group") + .required( + "block", + SyntaxShape::Block(Some(vec![SyntaxShape::Any])), + "the block to run on each group", + ) + .category(Category::Filters) + } + + fn usage(&self) -> &str { + "Runs a block on groups of `group_size` rows of a table at a time." + } + + fn examples(&self) -> Vec { + let stream_test_1 = vec![ + Value::Int { + val: 3, + span: Span::test_data(), + }, + Value::Int { + val: 7, + span: Span::test_data(), + }, + ]; + + vec![Example { + example: "echo [1 2 3 4] | each group 2 { $it.0 + $it.1 }", + description: "Multiplies elements in list", + result: Some(Value::List { + vals: stream_test_1, + span: Span::test_data(), + }), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + let group_size: Spanned = call.req(engine_state, stack, 0)?; + let capture_block: CaptureBlock = call.req(engine_state, stack, 1)?; + let ctrlc = engine_state.ctrlc.clone(); + + //FIXME: add in support for external redirection when engine-q supports it generally + + let each_group_iterator = EachGroupIterator { + block: capture_block, + engine_state: engine_state.clone(), + stack: stack.clone(), + group_size: group_size.item, + input: Box::new(input.into_iter()), + span: call.head, + }; + + Ok(each_group_iterator.flatten().into_pipeline_data(ctrlc)) + } +} + +struct EachGroupIterator { + block: CaptureBlock, + engine_state: EngineState, + stack: Stack, + group_size: usize, + input: Box + Send>, + span: Span, +} + +impl Iterator for EachGroupIterator { + type Item = PipelineData; + + fn next(&mut self) -> Option { + let mut group = vec![]; + let mut current_count = 0; + + loop { + let item = self.input.next(); + + match item { + Some(v) => { + group.push(v); + + current_count += 1; + if current_count >= self.group_size { + break; + } + } + None => break, + } + } + + if group.is_empty() { + return None; + } + + Some(run_block_on_vec( + group, + self.block.clone(), + self.engine_state.clone(), + self.stack.clone(), + self.span, + )) + } +} + +pub(crate) fn run_block_on_vec( + input: Vec, + capture_block: CaptureBlock, + engine_state: EngineState, + stack: Stack, + span: Span, +) -> PipelineData { + let value = Value::List { vals: input, span }; + + let mut stack = stack.captures_to_stack(&capture_block.captures); + + let block = engine_state.get_block(capture_block.block_id); + + if let Some(var) = block.signature.get_positional(0) { + if let Some(var_id) = &var.var_id { + stack.add_var(*var_id, value); + } + } + + match eval_block(&engine_state, &mut stack, block, PipelineData::new(span)) { + Ok(pipeline) => pipeline, + Err(error) => Value::Error { error }.into_pipeline_data(), + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_examples() { + use crate::test_examples; + + test_examples(EachGroup {}) + } +} diff --git a/crates/nu-command/src/filters/mod.rs b/crates/nu-command/src/filters/mod.rs index 7f5e14cb9d..adfe753083 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_group; mod empty; mod every; mod find; @@ -51,6 +52,7 @@ pub use compact::Compact; pub use default::Default; pub use drop::*; pub use each::Each; +pub use each_group::EachGroup; 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 eb38281600..988b79fd9f 100644 --- a/src/tests/test_engine.rs +++ b/src/tests/test_engine.rs @@ -72,7 +72,7 @@ fn in_variable_6() -> TestResult { #[test] fn help_works_with_missing_requirements() -> TestResult { - run_test(r#"each --help | lines | length"#, "16") + run_test(r#"each --help | lines | length"#, "19") } #[test]