diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs index 8743082ce2..671421e668 100644 --- a/crates/nu-command/src/default_context.rs +++ b/crates/nu-command/src/default_context.rs @@ -51,6 +51,7 @@ pub fn create_default_context() -> EngineState { Collect, Drop, DropColumn, + DropNth, Each, First, Get, diff --git a/crates/nu-command/src/filters/drop/mod.rs b/crates/nu-command/src/filters/drop/mod.rs index 9b9da9cebf..d3d59bd349 100644 --- a/crates/nu-command/src/filters/drop/mod.rs +++ b/crates/nu-command/src/filters/drop/mod.rs @@ -1,5 +1,7 @@ pub mod column; pub mod command; +pub mod nth; pub use column::DropColumn; pub use command::Drop; +pub use nth::DropNth; diff --git a/crates/nu-command/src/filters/drop/nth.rs b/crates/nu-command/src/filters/drop/nth.rs new file mode 100644 index 0000000000..f772dd162f --- /dev/null +++ b/crates/nu-command/src/filters/drop/nth.rs @@ -0,0 +1,122 @@ +use nu_engine::CallExt; +use nu_protocol::ast::Call; +use nu_protocol::engine::{Command, EngineState, Stack}; +use nu_protocol::{ + Category, Example, IntoInterruptiblePipelineData, PipelineData, PipelineIterator, ShellError, + Signature, Span, SyntaxShape, Value, +}; + +#[derive(Clone)] +pub struct DropNth; + +impl Command for DropNth { + fn name(&self) -> &str { + "drop nth" + } + + fn signature(&self) -> Signature { + Signature::build("drop nth") + .rest("rest", SyntaxShape::Int, "the number of the row to drop") + .category(Category::Filters) + } + + fn usage(&self) -> &str { + "Drop the selected rows." + } + + fn examples(&self) -> Vec { + vec![ + Example { + example: "[sam,sarah,2,3,4,5] | drop nth 0 1 2", + description: "Drop the first, second, and third row", + result: Some(Value::List { + vals: vec![Value::test_int(3), Value::test_int(4), Value::test_int(5)], + span: Span::unknown(), + }), + }, + Example { + example: "[0,1,2,3,4,5] | drop nth 0 1 2", + description: "Drop the first, second, and third row", + result: Some(Value::List { + vals: vec![Value::test_int(3), Value::test_int(4), Value::test_int(5)], + span: Span::unknown(), + }), + }, + Example { + example: "[0,1,2,3,4,5] | drop nth 0 2 4", + description: "Drop rows 0 2 4", + result: Some(Value::List { + vals: vec![Value::test_int(1), Value::test_int(3), Value::test_int(5)], + span: Span::unknown(), + }), + }, + Example { + example: "[0,1,2,3,4,5] | drop nth 2 0 4", + description: "Drop rows 2 0 4", + result: Some(Value::List { + vals: vec![Value::test_int(1), Value::test_int(3), Value::test_int(5)], + span: Span::unknown(), + }), + }, + ] + } + + fn run( + &self, + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, + ) -> Result { + let mut rows: Vec = call.rest(engine_state, stack, 0)?; + rows.sort_unstable(); + let pipeline_iter: PipelineIterator = input.into_iter(); + + Ok(DropNthIterator { + input: pipeline_iter, + rows, + current: 0, + } + .into_pipeline_data(engine_state.ctrlc.clone())) + } +} + +struct DropNthIterator { + input: PipelineIterator, + rows: Vec, + current: usize, +} + +impl Iterator for DropNthIterator { + type Item = Value; + + fn next(&mut self) -> Option { + loop { + if let Some(row) = self.rows.get(0) { + if self.current == *row { + self.rows.remove(0); + self.current += 1; + let _ = self.input.next(); + continue; + } else { + self.current += 1; + return self.input.next(); + } + } else { + return self.input.next(); + } + } + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_examples() { + use crate::test_examples; + + test_examples(DropNth {}) + } +}