From 463dd48180ac709585a6190907672da59ef4d6ed Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s=20N=2E=20Robalino?= <andres@androbtech.com>
Date: Fri, 13 Aug 2021 12:48:05 -0500
Subject: [PATCH] Flexible dropping of rows (by desired row number) (#3917)

We very well support `nth 0 2 3 --skip 1 4` to select particular rows and skip some using a flag. However, in practice we deal with tables (whether they come from parsing or loading files and whatnot) where we don't know the size of the table up front (and everytime we have these, they may have different sizes). There are also other use cases when we use intermediate tables during processing and wish to always drop certain rows and **keep the rest**.

Usage:

```
... | drop nth 0
... | drop nth 3 8
```
---
 .../src/commands/filters/drop/mod.rs          |  2 +
 .../src/commands/filters/drop/nth.rs          | 91 +++++++++++++++++++
 crates/nu-command/src/commands/filters/nth.rs | 13 ---
 crates/nu-command/src/commands/mod.rs         |  4 +
 crates/nu-command/src/default_context.rs      |  1 +
 5 files changed, 98 insertions(+), 13 deletions(-)
 create mode 100644 crates/nu-command/src/commands/filters/drop/nth.rs

diff --git a/crates/nu-command/src/commands/filters/drop/mod.rs b/crates/nu-command/src/commands/filters/drop/mod.rs
index 428affff8f..407d1c1006 100644
--- a/crates/nu-command/src/commands/filters/drop/mod.rs
+++ b/crates/nu-command/src/commands/filters/drop/mod.rs
@@ -1,5 +1,7 @@
 mod column;
 mod command;
+mod nth;
 
 pub use column::SubCommand as DropColumn;
 pub use command::Command as Drop;
+pub use nth::SubCommand as DropNth;
diff --git a/crates/nu-command/src/commands/filters/drop/nth.rs b/crates/nu-command/src/commands/filters/drop/nth.rs
new file mode 100644
index 0000000000..2208751703
--- /dev/null
+++ b/crates/nu-command/src/commands/filters/drop/nth.rs
@@ -0,0 +1,91 @@
+use crate::prelude::*;
+use nu_engine::WholeStreamCommand;
+use nu_errors::ShellError;
+use nu_protocol::{Signature, SyntaxShape, Value};
+use nu_source::Tagged;
+
+pub struct SubCommand;
+
+impl WholeStreamCommand for SubCommand {
+    fn name(&self) -> &str {
+        "drop nth"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build("drop nth")
+            .required(
+                "row number",
+                SyntaxShape::Int,
+                "the number of the row to drop",
+            )
+            .rest(SyntaxShape::Any, "Optionally drop more rows")
+    }
+
+    fn usage(&self) -> &str {
+        "Drops the selected rows."
+    }
+
+    fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
+        drop(args)
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![
+            Example {
+                description: "Drop the second row",
+                example: "echo [first second third] | drop nth 1",
+                result: Some(vec![Value::from("first"), Value::from("third")]),
+            },
+            Example {
+                description: "Drop the first and third rows",
+                example: "echo [first second third] | drop nth 0 2",
+                result: Some(vec![Value::from("second")]),
+            },
+        ]
+    }
+}
+
+fn drop(args: CommandArgs) -> Result<OutputStream, ShellError> {
+    let row_number: Tagged<u64> = args.req(0)?;
+    let and_rows: Vec<Tagged<u64>> = args.rest(1)?;
+    let input = args.input;
+
+    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();
+
+    Ok(DropNthIterator {
+        input,
+        rows,
+        current: 0,
+    }
+    .into_output_stream())
+}
+
+struct DropNthIterator {
+    input: InputStream,
+    rows: Vec<usize>,
+    current: usize,
+}
+
+impl Iterator for DropNthIterator {
+    type Item = Value;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        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();
+            }
+        }
+    }
+}
diff --git a/crates/nu-command/src/commands/filters/nth.rs b/crates/nu-command/src/commands/filters/nth.rs
index cbe31fa695..908026d247 100644
--- a/crates/nu-command/src/commands/filters/nth.rs
+++ b/crates/nu-command/src/commands/filters/nth.rs
@@ -112,16 +112,3 @@ impl Iterator for NthIterator {
         }
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::Nth;
-    use super::ShellError;
-
-    #[test]
-    fn examples_work_as_expected() -> Result<(), ShellError> {
-        use crate::examples::test as test_examples;
-
-        test_examples(Nth {})
-    }
-}
diff --git a/crates/nu-command/src/commands/mod.rs b/crates/nu-command/src/commands/mod.rs
index c2a8ac4aa5..41c7aa0f11 100644
--- a/crates/nu-command/src/commands/mod.rs
+++ b/crates/nu-command/src/commands/mod.rs
@@ -65,12 +65,16 @@ mod tests {
 
     fn full_tests() -> Vec<Command> {
         vec![
+            whole_stream_command(Drop),
+            whole_stream_command(DropNth),
+            whole_stream_command(DropColumn),
             whole_stream_command(Append),
             whole_stream_command(GroupBy),
             whole_stream_command(Insert),
             whole_stream_command(MoveColumn),
             whole_stream_command(Update),
             whole_stream_command(Empty),
+            whole_stream_command(Nth),
             // whole_stream_command(Select),
             // whole_stream_command(Get),
             // Str Command Suite
diff --git a/crates/nu-command/src/default_context.rs b/crates/nu-command/src/default_context.rs
index 795feb6d43..20c14d21c4 100644
--- a/crates/nu-command/src/default_context.rs
+++ b/crates/nu-command/src/default_context.rs
@@ -150,6 +150,7 @@ pub fn create_default_context(interactive: bool) -> Result<EvaluationContext, Bo
             whole_stream_command(Every),
             whole_stream_command(Nth),
             whole_stream_command(Drop),
+            whole_stream_command(DropNth),
             whole_stream_command(Format),
             whole_stream_command(FileSize),
             whole_stream_command(Where),