diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs index 6196e6a80a..8660b97ee4 100644 --- a/crates/nu-command/src/default_context.rs +++ b/crates/nu-command/src/default_context.rs @@ -55,6 +55,9 @@ pub fn create_default_context() -> EngineState { Each, First, Get, + Keep, + KeepUntil, + KeepWhile, Last, Length, Lines, diff --git a/crates/nu-command/src/filters/keep/command.rs b/crates/nu-command/src/filters/keep/command.rs new file mode 100644 index 0000000000..fd993bd6f0 --- /dev/null +++ b/crates/nu-command/src/filters/keep/command.rs @@ -0,0 +1,100 @@ +use std::convert::TryInto; + +use nu_engine::CallExt; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, IntoInterruptiblePipelineData, PipelineData, ShellError, Signature, Span, + SyntaxShape, Value, +}; + +#[derive(Clone)] +pub struct Keep; + +impl Command for Keep { + fn name(&self) -> &str { + "keep" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .optional("n", SyntaxShape::Int, "the number of elements to keep") + .category(Category::Filters) + } + + fn usage(&self) -> &str { + "Keep the first n elements of the input." + } + + fn examples(&self) -> Vec { + vec![ + Example { + description: "Keep two elements", + example: "echo [[editions]; [2015] [2018] [2021]] | keep 2", + result: Some(Value::List { + vals: vec![ + Value::Record { + cols: vec!["editions".to_owned()], + vals: vec![Value::from(2015)], + span: Span::unknown(), + }, + Value::Record { + cols: vec!["editions".to_owned()], + vals: vec![Value::from(2018)], + span: Span::unknown(), + }, + ], + span: Span::unknown(), + }), + }, + Example { + description: "Keep the first value", + example: "echo [2 4 6 8] | keep", + result: Some(Value::List { + vals: vec![Value::from(2)], + span: Span::unknown(), + }), + }, + ] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + let n: Option = call.opt(engine_state, stack, 0)?; + + let n: usize = match n { + Some(Value::Int { val, span }) => val.try_into().map_err(|err| { + ShellError::UnsupportedInput( + format!("Could not convert {} to unsigned integer: {}", val, err), + span, + ) + })?, + Some(_) => { + let span = call.head; + return Err(ShellError::TypeMismatch("expected integer".into(), span)); + } + None => 1, + }; + + let ctrlc = engine_state.ctrlc.clone(); + + Ok(input.into_iter().take(n).into_pipeline_data(ctrlc)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_examples() { + use crate::test_examples; + + test_examples(Keep {}) + } +} diff --git a/crates/nu-command/src/filters/keep/mod.rs b/crates/nu-command/src/filters/keep/mod.rs new file mode 100644 index 0000000000..681d472939 --- /dev/null +++ b/crates/nu-command/src/filters/keep/mod.rs @@ -0,0 +1,7 @@ +mod command; +mod until; +mod while_; + +pub use command::Keep; +pub use until::KeepUntil; +pub use while_::KeepWhile; diff --git a/crates/nu-command/src/filters/keep/until.rs b/crates/nu-command/src/filters/keep/until.rs new file mode 100644 index 0000000000..7c7c37a265 --- /dev/null +++ b/crates/nu-command/src/filters/keep/until.rs @@ -0,0 +1,90 @@ +use nu_engine::eval_block; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, IntoInterruptiblePipelineData, PipelineData, ShellError, Signature, Span, + SyntaxShape, Value, +}; + +#[derive(Clone)] +pub struct KeepUntil; + +impl Command for KeepUntil { + fn name(&self) -> &str { + "keep until" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .required( + "predicate", + SyntaxShape::RowCondition, + "the predicate that kept element must not match", + ) + .category(Category::Filters) + } + + fn usage(&self) -> &str { + "Keep elements of the input until a predicate is true." + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Keep until the element is positive", + example: "echo [-1 -2 9 1] | keep until $it > 0", + result: Some(Value::List { + vals: vec![Value::from(-1), Value::from(-2)], + span: Span::unknown(), + }), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + let span = call.head; + + let predicate = &call.positional[0]; + let block_id = predicate + .as_row_condition_block() + .ok_or_else(|| ShellError::TypeMismatch("expected row condition".to_owned(), span))?; + + let block = engine_state.get_block(block_id).clone(); + let var_id = block.signature.get_positional(0).and_then(|arg| arg.var_id); + + let mut stack = stack.collect_captures(&block.captures); + + let ctrlc = engine_state.ctrlc.clone(); + let engine_state = engine_state.clone(); + + Ok(input + .into_iter() + .take_while(move |value| { + if let Some(var_id) = var_id { + stack.add_var(var_id, value.clone()); + } + + !eval_block(&engine_state, &mut stack, &block, PipelineData::new(span)) + .map_or(false, |pipeline_data| { + pipeline_data.into_value(span).is_true() + }) + }) + .into_pipeline_data(ctrlc)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_examples() { + use crate::test_examples; + + test_examples(KeepUntil) + } +} diff --git a/crates/nu-command/src/filters/keep/while_.rs b/crates/nu-command/src/filters/keep/while_.rs new file mode 100644 index 0000000000..67a91c3891 --- /dev/null +++ b/crates/nu-command/src/filters/keep/while_.rs @@ -0,0 +1,90 @@ +use nu_engine::eval_block; +use nu_protocol::{ + ast::Call, + engine::{Command, EngineState, Stack}, + Category, Example, IntoInterruptiblePipelineData, PipelineData, ShellError, Signature, Span, + SyntaxShape, Value, +}; + +#[derive(Clone)] +pub struct KeepWhile; + +impl Command for KeepWhile { + fn name(&self) -> &str { + "keep while" + } + + fn signature(&self) -> Signature { + Signature::build(self.name()) + .required( + "predicate", + SyntaxShape::RowCondition, + "the predicate that kept element must not match", + ) + .category(Category::Filters) + } + + fn usage(&self) -> &str { + "Keep elements of the input while a predicate is true." + } + + fn examples(&self) -> Vec { + vec![Example { + description: "Keep while the element is negative", + example: "echo [-1 -2 9 1] | keep while $it < 0", + result: Some(Value::List { + vals: vec![Value::from(-1), Value::from(-2)], + span: Span::unknown(), + }), + }] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + let span = call.head; + + let predicate = &call.positional[0]; + let block_id = predicate + .as_row_condition_block() + .ok_or_else(|| ShellError::TypeMismatch("expected row condition".to_owned(), span))?; + + let block = engine_state.get_block(block_id).clone(); + let var_id = block.signature.get_positional(0).and_then(|arg| arg.var_id); + + let mut stack = stack.collect_captures(&block.captures); + + let ctrlc = engine_state.ctrlc.clone(); + let engine_state = engine_state.clone(); + + Ok(input + .into_iter() + .take_while(move |value| { + if let Some(var_id) = var_id { + stack.add_var(var_id, value.clone()); + } + + eval_block(&engine_state, &mut stack, &block, PipelineData::new(span)) + .map_or(false, |pipeline_data| { + pipeline_data.into_value(span).is_true() + }) + }) + .into_pipeline_data(ctrlc)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_examples() { + use crate::test_examples; + + test_examples(KeepWhile) + } +} diff --git a/crates/nu-command/src/filters/mod.rs b/crates/nu-command/src/filters/mod.rs index f2bf1cd345..1ebb2044b1 100644 --- a/crates/nu-command/src/filters/mod.rs +++ b/crates/nu-command/src/filters/mod.rs @@ -6,6 +6,7 @@ mod drop; mod each; mod first; mod get; +mod keep; mod last; mod length; mod lines; @@ -32,6 +33,7 @@ pub use drop::*; pub use each::Each; pub use first::First; pub use get::Get; +pub use keep::*; pub use last::Last; pub use length::Length; pub use lines::Lines;