diff --git a/crates/nu-command/src/commands/filters/drop/nth.rs b/crates/nu-command/src/commands/filters/drop/nth.rs index 2226c60aad..9e5b461e87 100644 --- a/crates/nu-command/src/commands/filters/drop/nth.rs +++ b/crates/nu-command/src/commands/filters/drop/nth.rs @@ -1,7 +1,8 @@ use crate::prelude::*; -use nu_engine::WholeStreamCommand; +use itertools::Either; +use nu_engine::{FromValue, WholeStreamCommand}; use nu_errors::ShellError; -use nu_protocol::{Signature, SyntaxShape, Value}; +use nu_protocol::{Range, Signature, SpannedTypeName, SyntaxShape, Value}; use nu_source::Tagged; pub struct SubCommand; @@ -14,11 +15,16 @@ impl WholeStreamCommand for SubCommand { fn signature(&self) -> Signature { Signature::build("drop nth") .required( - "row number", - SyntaxShape::Int, - "the number of the row to drop", + "row number or row range", + // FIXME: we can make this accept either Int or Range when we can compose SyntaxShapes + SyntaxShape::Any, + "the number of the row to drop or a range to drop consecutive rows", + ) + .rest( + "rest", + SyntaxShape::Any, + "Optionally drop more rows (Ignored if first argument is a range)", ) - .rest("rest", SyntaxShape::Any, "Optionally drop more rows") } fn usage(&self) -> &str { @@ -41,18 +47,46 @@ impl WholeStreamCommand for SubCommand { example: "echo [first second third] | drop nth 0 2", result: Some(vec![Value::from("second")]), }, + Example { + description: "Drop range rows from second to fourth", + example: "echo [first second third fourth fifth] | drop nth (1..3)", + result: Some(vec![Value::from("first"), Value::from("fifth")]), + }, ] } } -fn drop(args: CommandArgs) -> Result { - let row_number: Tagged = args.req(0)?; - let and_rows: Vec> = args.rest(1)?; - let input = args.input; +fn extract_int_or_range(args: &CommandArgs) -> Result, ShellError> { + let value = args.req::(0)?; - let mut rows: Vec<_> = and_rows.into_iter().map(|x| x.item as usize).collect(); - rows.push(row_number.item as usize); - rows.sort_unstable(); + let int_opt = value.as_u64().map(Either::Left).ok(); + let range_opt = FromValue::from_value(&value).map(Either::Right).ok(); + + int_opt + .or(range_opt) + .ok_or_else(|| ShellError::type_error("int or range", value.spanned_type_name())) +} + +fn drop(args: CommandArgs) -> Result { + let number_or_range = extract_int_or_range(&args)?; + let rows = match number_or_range { + Either::Left(row_number) => { + let and_rows: Vec> = args.rest(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(row_range) => { + let from = row_range.min_u64()? as usize; + let to = row_range.max_u64()? as usize; + + (from..=to).collect() + } + }; + + let input = args.input; Ok(DropNthIterator { input, diff --git a/crates/nu-command/tests/commands/drop.rs b/crates/nu-command/tests/commands/drop.rs index 613784b750..e723ba7690 100644 --- a/crates/nu-command/tests/commands/drop.rs +++ b/crates/nu-command/tests/commands/drop.rs @@ -66,3 +66,25 @@ fn more_rows_than_table_has() { assert_eq!(actual.out, "0"); } + +#[test] +fn nth_range_inclusive() { + let actual = nu!(cwd: ".", "echo 10..15 | drop nth (2..3) | to json"); + + assert_eq!(actual.out, "[10,11,14,15]"); +} + +#[test] +fn nth_range_exclusive() { + let actual = nu!(cwd: ".", "echo 10..15 | drop nth (1..<3) | to json"); + + assert_eq!(actual.out, "[10,13,14,15]"); +} + +#[test] +fn nth_missing_first_argument() { + let actual = nu!(cwd: ".", "echo 10..15 | drop nth \"\""); + + assert!(actual.err.contains("Expected int or range")); + assert!(actual.err.contains("found string")); +}