From d320ffe742b9a8bb418845947685579d6d8aff2e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9s=20N=2E=20Robalino?= <andres@androbtech.com>
Date: Mon, 25 Nov 2019 17:10:09 -0500
Subject: [PATCH] nth can select more than one row at a time.

---
 README.md              |  2 +-
 src/commands/nth.rs    | 50 ++++++++++++++++++++++++++++++++----------
 tests/commands_test.rs | 39 ++++++++++++++++++++++++++++++++
 3 files changed, 79 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index 5baa3e5ce0..1e804d22c1 100644
--- a/README.md
+++ b/README.md
@@ -281,7 +281,7 @@ Nu adheres closely to a set of goals that make up its design philosophy. As feat
 | inc (column-or-column-path) | Increment a value or version. Optionally use the column of a table |
 | insert column-or-column-path value | Insert a new column to the table |
 | last amount | Show only the last number of rows |
-| nth row-number | Return only the selected row |
+| nth ...row-numbers | Return only the selected rows |
 | pick ...columns | Down-select table to only these columns |
 | pivot --header-row <headers> | Pivot the tables, making columns into rows and vice versa |
 | prepend row-data | Prepend a row to the beginning of the table |
diff --git a/src/commands/nth.rs b/src/commands/nth.rs
index bcd3057879..6e495b25aa 100644
--- a/src/commands/nth.rs
+++ b/src/commands/nth.rs
@@ -5,7 +5,8 @@ use crate::prelude::*;
 
 #[derive(Deserialize)]
 struct NthArgs {
-    amount: Tagged<i64>,
+    row_number: Tagged<u64>,
+    rest: Vec<Tagged<u64>>,
 }
 
 pub struct Nth;
@@ -16,15 +17,17 @@ impl WholeStreamCommand for Nth {
     }
 
     fn signature(&self) -> Signature {
-        Signature::build("nth").required(
-            "row number",
-            SyntaxShape::Any,
-            "the number of the row to return",
-        )
+        Signature::build("nth")
+            .required(
+                "row number",
+                SyntaxShape::Any,
+                "the number of the row to return",
+            )
+            .rest(SyntaxShape::Any, "Optionally return more rows")
     }
 
     fn usage(&self) -> &str {
-        "Return only the selected row"
+        "Return only the selected rows"
     }
 
     fn run(
@@ -37,10 +40,35 @@ impl WholeStreamCommand for Nth {
 }
 
 fn nth(
-    NthArgs { amount }: NthArgs,
+    NthArgs {
+        row_number,
+        rest: and_rows,
+    }: NthArgs,
     RunnableContext { input, .. }: RunnableContext,
 ) -> Result<OutputStream, ShellError> {
-    Ok(OutputStream::from_input(
-        input.values.skip(amount.item as u64).take(1),
-    ))
+    let stream = input
+        .values
+        .enumerate()
+        .map(move |(idx, item)| {
+            let row_number = vec![row_number.clone()];
+
+            let row_numbers = vec![&row_number, &and_rows]
+                .into_iter()
+                .flatten()
+                .collect::<Vec<&Tagged<u64>>>();
+
+            let mut result = VecDeque::new();
+
+            if row_numbers
+                .iter()
+                .any(|requested| requested.item == idx as u64)
+            {
+                result.push_back(ReturnSuccess::value(item.clone()));
+            }
+
+            result
+        })
+        .flatten();
+
+    Ok(stream.to_output_stream())
 }
diff --git a/tests/commands_test.rs b/tests/commands_test.rs
index 59621539db..89ab430b36 100644
--- a/tests/commands_test.rs
+++ b/tests/commands_test.rs
@@ -3,6 +3,45 @@ mod helpers;
 use helpers as h;
 use helpers::{Playground, Stub::*};
 
+#[test]
+fn nth_selects_a_row() {
+    Playground::setup("nth_test_1", |dirs, sandbox| {
+        sandbox.with_files(vec![EmptyFile("notes.txt"), EmptyFile("arepas.txt")]);
+
+        let actual = nu!(
+            cwd: dirs.test(), h::pipeline(
+            r#"
+                ls
+                | sort-by name
+                | nth 0
+                | get name
+                | echo $it
+            "#
+        ));
+
+        assert_eq!(actual, "arepas.txt");
+    });
+}
+
+#[test]
+fn nth_selects_many_rows() {
+    Playground::setup("nth_test_2", |dirs, sandbox| {
+        sandbox.with_files(vec![EmptyFile("notes.txt"), EmptyFile("arepas.txt")]);
+
+        let actual = nu!(
+            cwd: dirs.test(), h::pipeline(
+            r#"
+                ls
+                | get name
+                | nth 1 0
+                | count
+                | echo $it
+            "#
+        ));
+
+        assert_eq!(actual, "2");
+    });
+}
 #[test]
 fn default_row_data_if_column_missing() {
     Playground::setup("default_test_1", |dirs, sandbox| {