diff --git a/crates/nu-command/src/filters/drop/nth.rs b/crates/nu-command/src/filters/drop/nth.rs index edcb8929..d3cd2372 100644 --- a/crates/nu-command/src/filters/drop/nth.rs +++ b/crates/nu-command/src/filters/drop/nth.rs @@ -1,9 +1,10 @@ +use itertools::Either; 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, + Category, Example, FromValue, IntoInterruptiblePipelineData, PipelineData, PipelineIterator, + Range, ShellError, Signature, Span, Spanned, SyntaxShape, Value, }; #[derive(Clone)] @@ -58,6 +59,14 @@ impl Command for DropNth { span: Span::test_data(), }), }, + Example { + description: "Drop range rows from second to fourth", + example: "echo [first second third fourth fifth] | drop nth (1..3)", + result: Some(Value::List { + vals: vec![Value::test_string("first"), Value::test_string("fifth")], + span: Span::test_data(), + }), + }, ] } @@ -68,12 +77,30 @@ impl Command for DropNth { 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(); + // let mut rows: Vec = call.rest(engine_state, stack, 0)?; + // rows.sort_unstable(); + // let pipeline_iter: PipelineIterator = input.into_iter(); + + let number_or_range = extract_int_or_range(engine_state, stack, call)?; + let rows = match number_or_range { + 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(row_range) => { + let from = row_range.from.as_integer()? as usize; + let to = row_range.to.as_integer()? as usize; + + (from..=to).collect() + } + }; Ok(DropNthIterator { - input: pipeline_iter, + input: input.into_iter(), rows, current: 0, } @@ -81,6 +108,26 @@ impl Command for DropNth { } } +fn extract_int_or_range( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, +) -> Result, ShellError> { + let value = call.req::(engine_state, stack, 0)?; + + let int_opt = value.as_integer().map(Either::Left).ok(); + let range_opt: Result = FromValue::from_value(&value); + + let range_opt = range_opt.map(Either::Right).ok(); + + int_opt.or(range_opt).ok_or_else(|| { + ShellError::TypeMismatch( + "int or range".into(), + value.span().unwrap_or(Span::new(0, 0)), + ) + }) +} + struct DropNthIterator { input: PipelineIterator, rows: Vec, diff --git a/crates/nu-command/tests/commands/drop.rs b/crates/nu-command/tests/commands/drop.rs index 3a2b9cdf..bb34a70e 100644 --- a/crates/nu-command/tests/commands/drop.rs +++ b/crates/nu-command/tests/commands/drop.rs @@ -70,3 +70,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")); +}