Merge main

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<PipelineData, ShellError> {
-        let mut rows: Vec<usize> = call.rest(engine_state, stack, 0)?;
-        rows.sort_unstable();
-        let pipeline_iter: PipelineIterator = input.into_iter();
+        // let mut rows: Vec<usize> = 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<Spanned<i64>> = 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<Either<i64, Range>, ShellError> {
+    let value = call.req::<Value>(engine_state, stack, 0)?;
+
+    let int_opt = value.as_integer().map(Either::Left).ok();
+    let range_opt: Result<nu_protocol::Range, ShellError> = 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<usize>,

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"));
+}