Simplify group/window into their own commands (#4760)

This commit is contained in:
JT 2022-03-06 20:01:29 -05:00 committed by GitHub
parent fbbbde1489
commit 4f43d75130
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 126 additions and 308 deletions

View file

@ -67,14 +67,13 @@ pub fn create_default_context(cwd: impl AsRef<Path>) -> EngineState {
DropColumn,
DropNth,
Each,
EachGroup,
EachWindow,
Empty,
Every,
Find,
First,
Flatten,
Get,
Group,
GroupBy,
Headers,
SplitBy,
@ -87,7 +86,6 @@ pub fn create_default_context(cwd: impl AsRef<Path>) -> EngineState {
Length,
Lines,
ParEach,
ParEachGroup,
Prepend,
Range,
Reduce,
@ -111,6 +109,7 @@ pub fn create_default_context(cwd: impl AsRef<Path>) -> EngineState {
Update,
UpdateCells,
Where,
Window,
Wrap,
Zip,
};

View file

@ -1,49 +1,62 @@
use nu_engine::{eval_block, CallExt};
use nu_engine::CallExt;
use nu_protocol::ast::Call;
use nu_protocol::engine::{CaptureBlock, Command, EngineState, Stack};
use nu_protocol::engine::{Command, EngineState, Stack};
use nu_protocol::{
Category, Example, IntoInterruptiblePipelineData, PipelineData, Signature, Span, Spanned,
SyntaxShape, Value,
};
#[derive(Clone)]
pub struct EachGroup;
pub struct Group;
impl Command for EachGroup {
impl Command for Group {
fn name(&self) -> &str {
"each group"
"group"
}
fn signature(&self) -> Signature {
Signature::build("each group")
Signature::build("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."
"Groups input into groups of `group_size`."
}
fn examples(&self) -> Vec<Example> {
let stream_test_1 = vec![
Value::Int {
val: 3,
Value::List {
vals: vec![
Value::Int {
val: 1,
span: Span::test_data(),
},
Value::Int {
val: 2,
span: Span::test_data(),
},
],
span: Span::test_data(),
},
Value::Int {
val: 7,
Value::List {
vals: vec![
Value::Int {
val: 3,
span: Span::test_data(),
},
Value::Int {
val: 4,
span: Span::test_data(),
},
],
span: Span::test_data(),
},
];
vec![Example {
example: "echo [1 2 3 4] | each group 2 { |it| $it.0 + $it.1 }",
description: "Echo the sum of each pair",
example: "echo [1 2 3 4] | group 2",
description: "Group the a list by pairs",
result: Some(Value::List {
vals: stream_test_1,
span: Span::test_data(),
@ -59,18 +72,12 @@ impl Command for EachGroup {
input: PipelineData,
) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
let group_size: Spanned<usize> = 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,
redirect_stdout: call.redirect_stdout,
redirect_stderr: call.redirect_stderr,
input: Box::new(input.into_iter()),
span: call.head,
};
@ -80,12 +87,7 @@ impl Command for EachGroup {
}
struct EachGroupIterator {
block: CaptureBlock,
engine_state: EngineState,
stack: Stack,
group_size: usize,
redirect_stdout: bool,
redirect_stderr: bool,
input: Box<dyn Iterator<Item = Value> + Send>,
span: Span,
}
@ -117,49 +119,10 @@ impl Iterator for EachGroupIterator {
return None;
}
Some(run_block_on_vec(
group,
self.block.clone(),
self.engine_state.clone(),
self.stack.clone(),
self.redirect_stdout,
self.redirect_stderr,
self.span,
))
}
}
pub(crate) fn run_block_on_vec(
input: Vec<Value>,
capture_block: CaptureBlock,
engine_state: EngineState,
stack: Stack,
redirect_stdout: bool,
redirect_stderr: bool,
span: Span,
) -> Value {
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),
redirect_stdout,
redirect_stderr,
) {
Ok(pipeline) => pipeline.into_value(span),
Err(error) => Value::Error { error },
Some(Value::List {
vals: group,
span: self.span,
})
}
}
@ -171,6 +134,6 @@ mod test {
fn test_examples() {
use crate::test_examples;
test_examples(EachGroup {})
test_examples(Group {})
}
}

View file

@ -7,14 +7,13 @@ mod compact;
mod default;
mod drop;
mod each;
mod each_group;
mod each_window;
mod empty;
mod every;
mod find;
mod first;
mod flatten;
mod get;
mod group;
mod group_by;
mod headers;
mod keep;
@ -24,7 +23,6 @@ mod lines;
mod merge;
mod move_;
mod par_each;
mod par_each_group;
mod prepend;
mod range;
mod reduce;
@ -43,6 +41,7 @@ mod uniq;
mod update;
mod update_cells;
mod where_;
mod window;
mod wrap;
mod zip_;
@ -55,14 +54,13 @@ pub use compact::Compact;
pub use default::Default;
pub use drop::*;
pub use each::Each;
pub use each_group::EachGroup;
pub use each_window::EachWindow;
pub use empty::Empty;
pub use every::Every;
pub use find::Find;
pub use first::First;
pub use flatten::Flatten;
pub use get::Get;
pub use group::Group;
pub use group_by::GroupBy;
pub use headers::Headers;
pub use keep::*;
@ -72,7 +70,6 @@ pub use lines::Lines;
pub use merge::Merge;
pub use move_::Move;
pub use par_each::ParEach;
pub use par_each_group::ParEachGroup;
pub use prepend::Prepend;
pub use range::Range;
pub use reduce::Reduce;
@ -91,5 +88,6 @@ pub use uniq::*;
pub use update::Update;
pub use update_cells::UpdateCells;
pub use where_::Where;
pub use window::Window;
pub use wrap::Wrap;
pub use zip_::Zip;

View file

@ -1,141 +0,0 @@
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, PipelineData, Signature, Spanned,
SyntaxShape, Value,
};
use rayon::prelude::*;
#[derive(Clone)]
pub struct ParEachGroup;
impl Command for ParEachGroup {
fn name(&self) -> &str {
"par-each group"
}
fn signature(&self) -> Signature {
Signature::build("par-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<Example> {
vec![Example {
example: "echo [1 2 3 4] | par-each group 2 {|it| $it.0 + $it.1 }",
description: "Multiplies elements in list",
result: None,
}]
}
fn run(
&self,
engine_state: &EngineState,
stack: &mut Stack,
call: &Call,
input: PipelineData,
) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
let group_size: Spanned<usize> = call.req(engine_state, stack, 0)?;
let capture_block: CaptureBlock = call.req(engine_state, stack, 1)?;
let ctrlc = engine_state.ctrlc.clone();
let span = call.head;
let redirect_stdout = call.redirect_stdout;
let redirect_stderr = call.redirect_stderr;
let stack = stack.captures_to_stack(&capture_block.captures);
//FIXME: add in support for external redirection when engine-q supports it generally
let each_group_iterator = EachGroupIterator {
group_size: group_size.item,
input: Box::new(input.into_iter()),
};
Ok(each_group_iterator
.par_bridge()
.map(move |x| {
let block = engine_state.get_block(capture_block.block_id);
let mut stack = stack.clone();
if let Some(var) = block.signature.get_positional(0) {
if let Some(var_id) = &var.var_id {
stack.add_var(*var_id, Value::List { vals: x, span });
}
}
match eval_block(
engine_state,
&mut stack,
block,
PipelineData::new(span),
redirect_stdout,
redirect_stderr,
) {
Ok(v) => v.into_value(span),
Err(error) => Value::Error { error },
}
})
.collect::<Vec<_>>()
.into_iter()
.into_pipeline_data(ctrlc))
}
}
struct EachGroupIterator {
group_size: usize,
input: Box<dyn Iterator<Item = Value> + Send>,
}
impl Iterator for EachGroupIterator {
type Item = Vec<Value>;
fn next(&mut self) -> Option<Self::Item> {
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(group)
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn test_examples() {
use crate::test_examples;
test_examples(ParEachGroup {})
}
}

View file

@ -1,21 +1,21 @@
use nu_engine::{eval_block, CallExt};
use nu_engine::CallExt;
use nu_protocol::ast::Call;
use nu_protocol::engine::{CaptureBlock, Command, EngineState, Stack};
use nu_protocol::engine::{Command, EngineState, Stack};
use nu_protocol::{
Category, Example, IntoInterruptiblePipelineData, PipelineData, Signature, Span, Spanned,
SyntaxShape, Value,
};
#[derive(Clone)]
pub struct EachWindow;
pub struct Window;
impl Command for EachWindow {
impl Command for Window {
fn name(&self) -> &str {
"each window"
"window"
}
fn signature(&self) -> Signature {
Signature::build("each window")
Signature::build("window")
.required("window_size", SyntaxShape::Int, "the size of each window")
.named(
"stride",
@ -23,52 +23,101 @@ impl Command for EachWindow {
"the number of rows to slide over between windows",
Some('s'),
)
.required(
"block",
SyntaxShape::Block(Some(vec![SyntaxShape::Any])),
"the block to run on each window",
)
.category(Category::Filters)
}
fn usage(&self) -> &str {
"Runs a block on window groups of `window_size` that slide by n rows."
"Creates a sliding window of `window_size` that slide by n rows/elements across input."
}
fn examples(&self) -> Vec<Example> {
let stream_test_1 = vec![
Value::Int {
val: 3,
Value::List {
vals: vec![
Value::Int {
val: 1,
span: Span::test_data(),
},
Value::Int {
val: 2,
span: Span::test_data(),
},
],
span: Span::test_data(),
},
Value::Int {
val: 5,
Value::List {
vals: vec![
Value::Int {
val: 2,
span: Span::test_data(),
},
Value::Int {
val: 3,
span: Span::test_data(),
},
],
span: Span::test_data(),
},
Value::Int {
val: 7,
Value::List {
vals: vec![
Value::Int {
val: 3,
span: Span::test_data(),
},
Value::Int {
val: 4,
span: Span::test_data(),
},
],
span: Span::test_data(),
},
];
let stream_test_2 = vec![
Value::Int {
val: 3,
Value::List {
vals: vec![
Value::Int {
val: 1,
span: Span::test_data(),
},
Value::Int {
val: 2,
span: Span::test_data(),
},
],
span: Span::test_data(),
},
Value::Int {
val: 9,
Value::List {
vals: vec![
Value::Int {
val: 4,
span: Span::test_data(),
},
Value::Int {
val: 5,
span: Span::test_data(),
},
],
span: Span::test_data(),
},
Value::Int {
val: 15,
Value::List {
vals: vec![
Value::Int {
val: 7,
span: Span::test_data(),
},
Value::Int {
val: 8,
span: Span::test_data(),
},
],
span: Span::test_data(),
},
];
vec![
Example {
example: "echo [1 2 3 4] | each window 2 { |it| $it.0 + $it.1 }",
example: "echo [1 2 3 4] | window 2",
description: "A sliding window of two elements",
result: Some(Value::List {
vals: stream_test_1,
@ -76,7 +125,7 @@ impl Command for EachWindow {
}),
},
Example {
example: "[1, 2, 3, 4, 5, 6, 7, 8] | each window 2 --stride 3 { |x| $x.0 + $x.1 }",
example: "[1, 2, 3, 4, 5, 6, 7, 8] | window 2 --stride 3",
description: "A sliding window of two elements, with a stride of 3",
result: Some(Value::List {
vals: stream_test_2,
@ -94,7 +143,6 @@ impl Command for EachWindow {
input: PipelineData,
) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
let group_size: Spanned<usize> = call.req(engine_state, stack, 0)?;
let capture_block: CaptureBlock = call.req(engine_state, stack, 1)?;
let ctrlc = engine_state.ctrlc.clone();
let stride: Option<usize> = call.get_flag(engine_state, stack, "stride")?;
@ -103,13 +151,8 @@ impl Command for EachWindow {
//FIXME: add in support for external redirection when engine-q supports it generally
let each_group_iterator = EachWindowIterator {
block: capture_block,
engine_state: engine_state.clone(),
stack: stack.clone(),
group_size: group_size.item,
input: Box::new(input.into_iter()),
redirect_stdout: call.redirect_stdout,
redirect_stderr: call.redirect_stderr,
span: call.head,
previous: vec![],
stride,
@ -120,13 +163,8 @@ impl Command for EachWindow {
}
struct EachWindowIterator {
block: CaptureBlock,
engine_state: EngineState,
stack: Stack,
group_size: usize,
input: Box<dyn Iterator<Item = Value> + Send>,
redirect_stdout: bool,
redirect_stderr: bool,
span: Span,
previous: Vec<Value>,
stride: usize,
@ -185,49 +223,10 @@ impl Iterator for EachWindowIterator {
self.previous = group.clone();
Some(run_block_on_vec(
group,
self.block.clone(),
self.engine_state.clone(),
self.stack.clone(),
self.redirect_stdout,
self.redirect_stderr,
self.span,
))
}
}
pub(crate) fn run_block_on_vec(
input: Vec<Value>,
capture_block: CaptureBlock,
engine_state: EngineState,
stack: Stack,
redirect_stdout: bool,
redirect_stderr: bool,
span: Span,
) -> Value {
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),
redirect_stdout,
redirect_stderr,
) {
Ok(pipeline) => pipeline.into_value(span),
Err(error) => Value::Error { error },
Some(Value::List {
vals: group,
span: self.span,
})
}
}
@ -239,6 +238,6 @@ mod test {
fn test_examples() {
use crate::test_examples;
test_examples(EachWindow {})
test_examples(Window {})
}
}

View file

@ -17,7 +17,7 @@ fn each_group_works() {
let actual = nu!(
cwd: "tests/fixtures/formats", pipeline(
r#"
echo [1 2 3 4 5 6] | each group 3 { |it| $it } | to json --raw
echo [1 2 3 4 5 6] | group 3 | to json --raw
"#
));
@ -29,7 +29,7 @@ fn each_window() {
let actual = nu!(
cwd: "tests/fixtures/formats", pipeline(
r#"
echo [1 2 3 4] | each window 3 { |it| $it } | to json --raw
echo [1 2 3 4] | window 3 | to json --raw
"#
));
@ -41,7 +41,7 @@ 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 { |it| echo $it } | to json --raw
echo [1 2 3 4 5 6] | window 3 -s 2 | to json --raw
"#
));

View file

@ -72,7 +72,7 @@ fn in_variable_6() -> TestResult {
#[test]
fn help_works_with_missing_requirements() -> TestResult {
run_test(r#"each --help | lines | length"#, "33")
run_test(r#"each --help | lines | length"#, "29")
}
#[test]