diff --git a/crates/nu-command/src/filters/drop/nth.rs b/crates/nu-command/src/filters/drop/nth.rs index ed4e0842d0..ebb3efaca5 100644 --- a/crates/nu-command/src/filters/drop/nth.rs +++ b/crates/nu-command/src/filters/drop/nth.rs @@ -1,6 +1,6 @@ -use itertools::Either; use nu_engine::command_prelude::*; use nu_protocol::{PipelineIterator, Range}; +use std::collections::VecDeque; use std::ops::Bound; #[derive(Clone)] @@ -18,13 +18,11 @@ impl Command for DropNth { (Type::list(Type::Any), Type::list(Type::Any)), ]) .allow_variants_without_examples(true) - .required( - "row number or row range", - // FIXME: we can make this accept either Int or Range when we can compose SyntaxShapes + .rest( + "rest", SyntaxShape::Any, - "The number of the row to drop or a range to drop consecutive rows.", + "The row numbers or ranges to drop.", ) - .rest("rest", SyntaxShape::Any, "The number of the row to drop.") .category(Category::Filters) } @@ -103,110 +101,125 @@ impl Command for DropNth { ) -> Result { let head = call.head; let metadata = input.metadata(); - let number_or_range = extract_int_or_range(engine_state, stack, call)?; - let rows = match number_or_range.item { - Either::Left(row_number) => { - let and_rows: Vec> = call.rest(engine_state, stack, 1)?; - let mut rows: Vec<_> = and_rows.into_iter().map(|x| x.item as usize).collect(); - rows.push(row_number as usize); - rows.sort_unstable(); - rows - } - Either::Right(Range::FloatRange(_)) => { - return Err(ShellError::UnsupportedInput { - msg: "float range".into(), - input: "value originates from here".into(), - msg_span: head, - input_span: number_or_range.span, - }); - } - Either::Right(Range::IntRange(range)) => { - // check for negative range inputs, e.g., (2..-5) - let end_negative = match range.end() { - Bound::Included(end) | Bound::Excluded(end) => end < 0, - Bound::Unbounded => false, - }; - if range.start().is_negative() || end_negative { - return Err(ShellError::UnsupportedInput { - msg: "drop nth accepts only positive ints".into(), - input: "value originates from here".into(), - msg_span: head, - input_span: number_or_range.span, - }); - } - // check if the upper bound is smaller than the lower bound, e.g., do not accept 4..2 - if range.step() < 0 { - return Err(ShellError::UnsupportedInput { - msg: "The upper bound needs to be equal or larger to the lower bound" - .into(), - input: "value originates from here".into(), - msg_span: head, - input_span: number_or_range.span, - }); - } + let args: Vec = call.rest(engine_state, stack, 0)?; + if args.is_empty() { + return Ok(input); + } - let start = range.start() as usize; + let (rows_to_drop, min_unbounded_start) = get_rows_to_drop(&args, head)?; - let end = match range.end() { - Bound::Included(end) => end as usize, - Bound::Excluded(end) => (end - 1) as usize, - Bound::Unbounded => { - return Ok(input - .into_iter() - .take(start) - .into_pipeline_data_with_metadata( - head, - engine_state.signals().clone(), - metadata, - )); - } - }; - - let end = if let PipelineData::Value(Value::List { vals, .. }, _) = &input { - end.min(vals.len() - 1) - } else { - end - }; - - (start..=end).collect() - } + let input = if let Some(cutoff) = min_unbounded_start { + input + .into_iter() + .take(cutoff) + .into_pipeline_data_with_metadata( + head, + engine_state.signals().clone(), + metadata.clone(), + ) + } else { + input }; Ok(DropNthIterator { input: input.into_iter(), - rows, + rows: rows_to_drop, current: 0, } .into_pipeline_data_with_metadata(head, engine_state.signals().clone(), metadata)) } } -fn extract_int_or_range( - engine_state: &EngineState, - stack: &mut Stack, - call: &Call, -) -> Result>, ShellError> { - let value: Value = call.req(engine_state, stack, 0)?; +fn get_rows_to_drop( + args: &[Value], + head: Span, +) -> Result<(VecDeque, Option), ShellError> { + let mut rows_to_drop = Vec::new(); + let mut min_unbounded_start: Option = None; - let int_opt = value.as_int().map(Either::Left).ok(); - let range_opt = value.as_range().map(Either::Right).ok(); + for value in args { + if let Ok(i) = value.as_int() { + if i < 0 { + return Err(ShellError::UnsupportedInput { + msg: "drop nth accepts only positive ints".into(), + input: "value originates from here".into(), + msg_span: head, + input_span: value.span(), + }); + } + rows_to_drop.push(i as usize); + } else if let Ok(range) = value.as_range() { + match range { + Range::IntRange(range) => { + let start = range.start(); + if start < 0 { + return Err(ShellError::UnsupportedInput { + msg: "drop nth accepts only positive ints".into(), + input: "value originates from here".into(), + msg_span: head, + input_span: value.span(), + }); + } - int_opt - .or(range_opt) - .ok_or_else(|| ShellError::TypeMismatch { - err_message: "int or range".into(), - span: value.span(), - }) - .map(|either| Spanned { - item: either, - span: value.span(), - }) + match range.end() { + Bound::Included(end) => { + if end < start { + return Err(ShellError::UnsupportedInput { + msg: "The upper bound must be greater than or equal to the lower bound".into(), + input: "value originates from here".into(), + msg_span: head, + input_span: value.span(), + }); + } + rows_to_drop.extend((start as usize)..=(end as usize)); + } + Bound::Excluded(end) => { + if end <= start { + return Err(ShellError::UnsupportedInput { + msg: "The upper bound must be greater than the lower bound" + .into(), + input: "value originates from here".into(), + msg_span: head, + input_span: value.span(), + }); + } + rows_to_drop.extend((start as usize)..(end as usize)); + } + Bound::Unbounded => { + let start_usize = start as usize; + min_unbounded_start = Some( + min_unbounded_start.map_or(start_usize, |s| s.min(start_usize)), + ); + } + } + } + Range::FloatRange(_) => { + return Err(ShellError::UnsupportedInput { + msg: "float range not supported".into(), + input: "value originates from here".into(), + msg_span: head, + input_span: value.span(), + }); + } + } + } else { + return Err(ShellError::TypeMismatch { + err_message: "Expected int or range".into(), + span: value.span(), + }); + } + } + + rows_to_drop.sort_unstable(); + rows_to_drop.dedup(); + + Ok((VecDeque::from(rows_to_drop), min_unbounded_start)) } struct DropNthIterator { input: PipelineIterator, - rows: Vec, + rows: VecDeque, current: usize, } @@ -215,9 +228,9 @@ impl Iterator for DropNthIterator { fn next(&mut self) -> Option { loop { - if let Some(row) = self.rows.first() { + if let Some(row) = self.rows.front() { if self.current == *row { - self.rows.remove(0); + self.rows.pop_front(); self.current += 1; let _ = self.input.next(); continue; diff --git a/crates/nu-command/tests/commands/drop.rs b/crates/nu-command/tests/commands/drop.rs index 2f7fb9eaaa..f9560def7c 100644 --- a/crates/nu-command/tests/commands/drop.rs +++ b/crates/nu-command/tests/commands/drop.rs @@ -135,3 +135,75 @@ fn more_columns_than_record_has() { assert_eq!(actual.out, "{}"); } + +#[test] +fn drop_single_index() { + let actual = nu!("echo 10..15 | drop nth 2 | to json --raw"); + assert_eq!(actual.out, "[10,11,13,14,15]"); +} + +#[test] +fn drop_multiple_indices() { + let actual = nu!("echo 0..10 | drop nth 1 3 | to json --raw"); + assert_eq!(actual.out, "[0,2,4,5,6,7,8,9,10]"); +} + +#[test] +fn drop_inclusive_range() { + let actual = nu!("echo 10..15 | drop nth (2..4) | to json --raw"); + assert_eq!(actual.out, "[10,11,15]"); +} + +#[test] +fn drop_exclusive_range() { + let actual = nu!("echo 10..15 | drop nth (2..<4) | to json --raw"); + assert_eq!(actual.out, "[10,11,14,15]"); +} + +#[test] +fn drop_unbounded_range() { + let actual = nu!("echo 0..5 | drop nth 3.. | to json --raw"); + assert_eq!(actual.out, "[0,1,2]"); +} + +#[test] +fn drop_multiple_ranges_including_unbounded() { + let actual = nu!(pipeline( + r#" + 0..30 + | drop nth 0..10 20.. + | to json --raw + "# + )); + + assert_eq!(actual.out, "[11,12,13,14,15,16,17,18,19]"); +} + +#[test] +fn drop_combination_of_unbounded_range_and_single_index() { + let actual = nu!(pipeline( + r#" + echo 0..15 + | drop nth 10.. 5 + | to json --raw + "# + )); + + assert_eq!(actual.out, "[0,1,2,3,4,6,7,8,9]"); +} + +#[test] +fn drop_combination_of_two_unbounded_ranges() { + let actual = nu!(pipeline( + r#" + echo 0..150 + | drop nth 0..100 999.. + | to json --raw + "# + )); + + let expected: Vec = (101..=150).collect(); + let expected_json = serde_json::to_string(&expected).unwrap(); + + assert_eq!(actual.out, expected_json); +}