From 2c5939dc7d7489a8ce6a07ca035298f131de2c77 Mon Sep 17 00:00:00 2001 From: Ritoban Roy-Chowdhury Date: Sun, 6 Sep 2020 22:54:52 -0700 Subject: [PATCH] `each group` and `each window` subcommands. (#2508) * First commit updating `config` to use subcommands (#2119) - Implemented `get` subcommand * Implmented `config set` as a subcommand. * Implemented `config set_into` as subcommand * Fixed base `config` command - Instead of outputting help, it now outputs the list of all configuration parameters. * Added `config clear` subcommand * Added `config load` and `config remove` subcommands * Added `config path` subcommand * fixed clippy * initial commit for implementing groups * each group works * each group is slightly cleaner + added example * Added `each window` subcommand - No support for stride flag yet * each window stride implemented * Added tests and minor documentation changes * fixed clippy * fixed clippy again --- crates/nu-cli/src/cli.rs | 2 + crates/nu-cli/src/commands.rs | 2 + .../src/commands/{each.rs => each/command.rs} | 0 crates/nu-cli/src/commands/each/group.rs | 133 ++++++++++++++++++ crates/nu-cli/src/commands/each/mod.rs | 9 ++ crates/nu-cli/src/commands/each/window.rs | 113 +++++++++++++++ crates/nu-cli/tests/commands/each.rs | 36 +++++ 7 files changed, 295 insertions(+) rename crates/nu-cli/src/commands/{each.rs => each/command.rs} (100%) create mode 100644 crates/nu-cli/src/commands/each/group.rs create mode 100644 crates/nu-cli/src/commands/each/mod.rs create mode 100644 crates/nu-cli/src/commands/each/window.rs diff --git a/crates/nu-cli/src/cli.rs b/crates/nu-cli/src/cli.rs index 1e49e8c55..ed445cbdb 100644 --- a/crates/nu-cli/src/cli.rs +++ b/crates/nu-cli/src/cli.rs @@ -205,6 +205,8 @@ pub fn create_default_context( whole_stream_command(Rename), whole_stream_command(Uniq), whole_stream_command(Each), + whole_stream_command(EachGroup), + whole_stream_command(EachWindow), whole_stream_command(IsEmpty), // Table manipulation whole_stream_command(Move), diff --git a/crates/nu-cli/src/commands.rs b/crates/nu-cli/src/commands.rs index f5398683c..e9743f534 100644 --- a/crates/nu-cli/src/commands.rs +++ b/crates/nu-cli/src/commands.rs @@ -155,6 +155,8 @@ pub(crate) use do_::Do; pub(crate) use drop::Drop; pub(crate) use du::Du; pub(crate) use each::Each; +pub(crate) use each::EachGroup; +pub(crate) use each::EachWindow; pub(crate) use echo::Echo; pub(crate) use if_::If; pub(crate) use is_empty::IsEmpty; diff --git a/crates/nu-cli/src/commands/each.rs b/crates/nu-cli/src/commands/each/command.rs similarity index 100% rename from crates/nu-cli/src/commands/each.rs rename to crates/nu-cli/src/commands/each/command.rs diff --git a/crates/nu-cli/src/commands/each/group.rs b/crates/nu-cli/src/commands/each/group.rs new file mode 100644 index 000000000..cb3678c6c --- /dev/null +++ b/crates/nu-cli/src/commands/each/group.rs @@ -0,0 +1,133 @@ +use crate::commands::each::process_row; +use crate::commands::WholeStreamCommand; +use crate::prelude::*; +use nu_errors::ShellError; +use nu_protocol::{ + hir::Block, hir::SpannedExpression, ReturnSuccess, Scope, Signature, SyntaxShape, + UntaggedValue, Value, +}; +use nu_source::Tagged; +use serde::Deserialize; + +pub struct EachGroup; + +#[derive(Deserialize)] +pub struct EachGroupArgs { + group_size: Tagged, + block: Block, + //numbered: Tagged, +} + +#[async_trait] +impl WholeStreamCommand 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, + "the block to run on each group", + ) + } + + fn usage(&self) -> &str { + "Runs a block on groups of `group_size` rows of a table at a time." + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Echo the sum of each pair", + example: "echo [1 2 3 4] | each group 2 { echo $it | math sum }", + result: None, + }] + } + + async fn run( + &self, + raw_args: CommandArgs, + registry: &CommandRegistry, + ) -> Result { + let registry = registry.clone(); + let head = Arc::new(raw_args.call_info.args.head.clone()); + let scope = Arc::new(raw_args.call_info.scope.clone()); + let context = Arc::new(Context::from_raw(&raw_args, ®istry)); + let (each_args, input): (EachGroupArgs, _) = raw_args.process(®istry).await?; + let block = Arc::new(each_args.block); + + Ok(input + .chunks(each_args.group_size.item) + .then(move |input| { + run_block_on_vec( + input, + block.clone(), + scope.clone(), + head.clone(), + context.clone(), + ) + }) + .flatten() + .to_output_stream()) + } +} + +pub(crate) fn run_block_on_vec( + input: Vec, + block: Arc, + scope: Arc, + head: Arc>, + context: Arc, +) -> impl Future { + let value = Value { + value: UntaggedValue::Table(input), + tag: Tag::unknown(), + }; + + async { + match process_row(block, scope, head, context, value).await { + Ok(s) => { + // We need to handle this differently depending on whether process_row + // returned just 1 value or if it returned multiple as a stream. + let vec = s.collect::>().await; + + // If it returned just one value, just take that value + if vec.len() == 1 { + return OutputStream::one(vec.into_iter().next().expect( + "This should be impossible, we just checked that vec.len() == 1.", + )); + } + + // If it returned multiple values, we need to put them into a table and + // return that. + let result = vec.into_iter().collect::, _>>(); + let result_table = match result { + Ok(t) => t, + Err(e) => return OutputStream::one(Err(e)), + }; + + let table = result_table + .into_iter() + .filter_map(|x| x.raw_value()) + .collect(); + + OutputStream::one(Ok(ReturnSuccess::Value(UntaggedValue::Table(table).into()))) + } + Err(e) => OutputStream::one(Err(e)), + } + } +} + +#[cfg(test)] +mod tests { + use super::EachGroup; + + #[test] + fn examples_work_as_expected() { + use crate::examples::test as test_examples; + + test_examples(EachGroup {}) + } +} diff --git a/crates/nu-cli/src/commands/each/mod.rs b/crates/nu-cli/src/commands/each/mod.rs new file mode 100644 index 000000000..cd656685f --- /dev/null +++ b/crates/nu-cli/src/commands/each/mod.rs @@ -0,0 +1,9 @@ +pub mod command; +pub mod group; +pub mod window; + +pub(crate) use command::make_indexed_item; +pub use command::process_row; +pub use command::Each; +pub use group::EachGroup; +pub use window::EachWindow; diff --git a/crates/nu-cli/src/commands/each/window.rs b/crates/nu-cli/src/commands/each/window.rs new file mode 100644 index 000000000..a633f39db --- /dev/null +++ b/crates/nu-cli/src/commands/each/window.rs @@ -0,0 +1,113 @@ +use crate::commands::each::group::run_block_on_vec; +use crate::commands::WholeStreamCommand; +use crate::prelude::*; +//use itertools::Itertools; +use nu_errors::ShellError; +use nu_protocol::{hir::Block, Primitive, Signature, SyntaxShape, UntaggedValue}; +use nu_source::Tagged; +use serde::Deserialize; + +pub struct EachWindow; + +#[derive(Deserialize)] +pub struct EachWindowArgs { + window_size: Tagged, + block: Block, + stride: Option>, +} + +#[async_trait] +impl WholeStreamCommand for EachWindow { + fn name(&self) -> &str { + "each window" + } + + fn signature(&self) -> Signature { + Signature::build("each window") + .required("window_size", SyntaxShape::Int, "the size of each window") + .named( + "stride", + SyntaxShape::Int, + "the number of rows to slide over between windows", + Some('s'), + ) + .required( + "block", + SyntaxShape::Block, + "the block to run on each group", + ) + } + + fn usage(&self) -> &str { + "Runs a block on sliding windows of `window_size` rows of a table at a time." + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Echo the sum of each window", + example: "echo [1 2 3 4] | each window 2 { echo $it | math sum }", + result: None, + }] + } + + async fn run( + &self, + raw_args: CommandArgs, + registry: &CommandRegistry, + ) -> Result { + let registry = registry.clone(); + let head = Arc::new(raw_args.call_info.args.head.clone()); + let scope = Arc::new(raw_args.call_info.scope.clone()); + let context = Arc::new(Context::from_raw(&raw_args, ®istry)); + let (each_args, mut input): (EachWindowArgs, _) = raw_args.process(®istry).await?; + let block = Arc::new(each_args.block); + + let mut window: Vec<_> = input + .by_ref() + .take(*each_args.window_size - 1) + .collect::>() + .await; + + // `window` must start with dummy values, which will be removed on the first iteration + let stride = each_args.stride.map(|x| *x).unwrap_or(1); + window.insert(0, UntaggedValue::Primitive(Primitive::Nothing).into()); + + Ok(input + .enumerate() + .then(move |(i, input)| { + // This would probably be more efficient if `last` was a VecDeque + // But we can't have that because it needs to be put into a Table + window.remove(0); + window.push(input); + + let block = block.clone(); + let scope = scope.clone(); + let head = head.clone(); + let context = context.clone(); + let local_window = window.clone(); + + async move { + if i % stride == 0 { + Some(run_block_on_vec(local_window, block, scope, head, context).await) + } else { + None + } + } + }) + .filter_map(|x| async { x }) + .flatten() + .to_output_stream()) + } +} + +#[cfg(test)] +mod tests { + use super::EachWindow; + + #[test] + fn examples_work_as_expected() { + use crate::examples::test as test_examples; + + test_examples(EachWindow {}) + } +} diff --git a/crates/nu-cli/tests/commands/each.rs b/crates/nu-cli/tests/commands/each.rs index f12ffe3b6..466fafa5f 100644 --- a/crates/nu-cli/tests/commands/each.rs +++ b/crates/nu-cli/tests/commands/each.rs @@ -11,3 +11,39 @@ fn each_works_separately() { assert_eq!(actual.out, "[11,12,13]"); } + +#[test] +fn each_group_works() { + let actual = nu!( + cwd: "tests/fixtures/formats", pipeline( + r#" + echo [1 2 3 4 5 6] | each group 3 { echo $it } | to json + "# + )); + + assert_eq!(actual.out, "[[1,2,3],[4,5,6]]"); +} + +#[test] +fn each_window() { + let actual = nu!( + cwd: "tests/fixtures/formats", pipeline( + r#" + echo [1 2 3 4] | each window 3 { echo $it } | to json + "# + )); + + assert_eq!(actual.out, "[[1,2,3],[2,3,4]]"); +} + +#[test] +fn each_window_stride() { + let actual = nu!( + cwd: "tests/fixtures/formats", pipeline( + r#" + echo [1 2 3 4 5 6] | each window 3 -s 2 { echo $it } | to json + "# + )); + + assert_eq!(actual.out, "[[1,2,3],[3,4,5]]"); +}