From 85587c0c2a21b66ac08a3af88a72f6eb5307e04a Mon Sep 17 00:00:00 2001
From: WindSoilder <WindSoilder@outlook.com>
Date: Wed, 9 Nov 2022 18:32:16 +0800
Subject: [PATCH] make `take` behave like `first` (#6893)

* fix take_1 behavior

* fix test case

* simplify code

* reverse back for first command

* fix example

* make arg required

* add test

* fix test message
---
 crates/nu-command/src/filters/take/take_.rs   | 34 +++++++------------
 crates/nu-command/tests/commands/take/rows.rs | 12 +++++++
 2 files changed, 24 insertions(+), 22 deletions(-)

diff --git a/crates/nu-command/src/filters/take/take_.rs b/crates/nu-command/src/filters/take/take_.rs
index d887f473b4..a99d651ca2 100644
--- a/crates/nu-command/src/filters/take/take_.rs
+++ b/crates/nu-command/src/filters/take/take_.rs
@@ -16,7 +16,7 @@ impl Command for Take {
 
     fn signature(&self) -> Signature {
         Signature::build("take")
-            .optional(
+            .required(
                 "n",
                 SyntaxShape::Int,
                 "starting from the front, the number of elements to return",
@@ -42,8 +42,11 @@ impl Command for Take {
         vec![
             Example {
                 description: "Return the first item of a list/table",
-                example: "[1 2 3] | take",
-                result: Some(Value::test_int(1)),
+                example: "[1 2 3] | take 1",
+                result: Some(Value::List {
+                    vals: vec![Value::test_int(1)],
+                    span: Span::test_data(),
+                }),
             },
             Example {
                 description: "Return the first 2 items of a list/table",
@@ -64,11 +67,7 @@ fn first_helper(
     input: PipelineData,
 ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
     let head = call.head;
-    let rows: Option<i64> = call.opt(engine_state, stack, 0)?;
-    let mut rows_desired: usize = match rows {
-        Some(x) => x as usize,
-        None => 1,
-    };
+    let mut rows_desired: usize = call.req(engine_state, stack, 0)?;
 
     let ctrlc = engine_state.ctrlc.clone();
     let metadata = input.metadata();
@@ -137,20 +136,11 @@ fn first_helper(
                         .set_metadata(metadata)),
                 }
             }
-            _ => {
-                if rows_desired == 1 {
-                    match input_peek.next() {
-                        Some(val) => Ok(val.into_pipeline_data()),
-                        None => Err(ShellError::AccessBeyondEndOfStream(head)),
-                    }
-                } else {
-                    Ok(input_peek
-                        .into_iter()
-                        .take(rows_desired)
-                        .into_pipeline_data(ctrlc)
-                        .set_metadata(metadata))
-                }
-            }
+            _ => Ok(input_peek
+                .into_iter()
+                .take(rows_desired)
+                .into_pipeline_data(ctrlc)
+                .set_metadata(metadata)),
         }
     } else {
         Err(ShellError::UnsupportedInput(
diff --git a/crates/nu-command/tests/commands/take/rows.rs b/crates/nu-command/tests/commands/take/rows.rs
index a4b1a78f2d..3771781d1e 100644
--- a/crates/nu-command/tests/commands/take/rows.rs
+++ b/crates/nu-command/tests/commands/take/rows.rs
@@ -29,3 +29,15 @@ fn rows() {
         assert_eq!(actual.out, "4");
     })
 }
+
+#[test]
+fn rows_with_no_arguments_should_lead_to_error() {
+    Playground::setup("take_test_2", |dirs, _sandbox| {
+        let actual = nu!(
+            cwd: dirs.test(), pipeline(
+            r#"[1 2 3] | take"#
+        ));
+
+        assert!(actual.err.contains("missing_positional"));
+    })
+}