Make take work like first (#5942)

crates/nu-command/src/filters/take/take_.rs

@@ -1,11 +1,9 @@
-use std::convert::TryInto;
-
 use nu_engine::CallExt;
+use nu_protocol::ast::Call;
+use nu_protocol::engine::{Command, EngineState, Stack};
 use nu_protocol::{
-    ast::Call,
+    Category, Example, IntoInterruptiblePipelineData, IntoPipelineData, PipelineData, ShellError,
-    engine::{Command, EngineState, Stack},
+    Signature, Span, SyntaxShape, Type, Value,
-    Category, Example, IntoInterruptiblePipelineData, PipelineData, ShellError, Signature, Span,
-    SyntaxShape, Value,
 };
 
 #[derive(Clone)]
@@ -17,45 +15,17 @@ impl Command for Take {
     }
 
     fn signature(&self) -> Signature {
-        Signature::build(self.name())
+        Signature::build("take")
-            .optional("n", SyntaxShape::Int, "the number of elements to take")
+            .optional(
+                "n",
+                SyntaxShape::Int,
+                "starting from the front, the number of elements to return",
+            )
             .category(Category::Filters)
     }
 
     fn usage(&self) -> &str {
-        "Take the first n elements of the input."
+        "Take only the first n elements."
-    }
-
-    fn examples(&self) -> Vec<Example> {
-        vec![
-            Example {
-                description: "Take two elements",
-                example: "echo [[editions]; [2015] [2018] [2021]] | take 2",
-                result: Some(Value::List {
-                    vals: vec![
-                        Value::Record {
-                            cols: vec!["editions".to_owned()],
-                            vals: vec![Value::test_int(2015)],
-                            span: Span::test_data(),
-                        },
-                        Value::Record {
-                            cols: vec!["editions".to_owned()],
-                            vals: vec![Value::test_int(2018)],
-                            span: Span::test_data(),
-                        },
-                    ],
-                    span: Span::test_data(),
-                }),
-            },
-            Example {
-                description: "Take the first value",
-                example: "echo [2 4 6 8] | take",
-                result: Some(Value::List {
-                    vals: vec![Value::test_int(2)],
-                    span: Span::test_data(),
-                }),
-            },
-        ]
     }
 
     fn run(
@@ -64,38 +34,134 @@ impl Command for Take {
         stack: &mut Stack,
         call: &Call,
         input: PipelineData,
-    ) -> Result<PipelineData, ShellError> {
+    ) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
-        let n: Option<Value> = call.opt(engine_state, stack, 0)?;
+        first_helper(engine_state, stack, call, input)
-        let metadata = input.metadata();
+    }
 
-        let n: usize = match n {
+    fn examples(&self) -> Vec<Example> {
-            Some(Value::Int { val, span }) => val.try_into().map_err(|err| {
+        vec![
-                ShellError::UnsupportedInput(
+            Example {
-                    format!("Could not convert {} to unsigned integer: {}", val, err),
+                description: "Return the first item of a list/table",
-                    span,
+                example: "[1 2 3] | take",
-                )
+                result: Some(Value::test_int(1)),
-            })?,
+            },
-            Some(_) => {
+            Example {
-                let span = call.head;
+                description: "Return the first 2 items of a list/table",
-                return Err(ShellError::TypeMismatch("expected integer".into(), span));
+                example: "[1 2 3] | take 2",
-            }
+                result: Some(Value::List {
-            None => 1,
+                    vals: vec![Value::test_int(1), Value::test_int(2)],
-        };
+                    span: Span::test_data(),
-
+                }),
-        let ctrlc = engine_state.ctrlc.clone();
+            },
-
+        ]
-        Ok(input
-            .into_iter()
-            .take(n)
-            .into_pipeline_data(ctrlc)
-            .set_metadata(metadata))
     }
 }
 
-#[cfg(test)]
+fn first_helper(
-mod tests {
+    engine_state: &EngineState,
-    use crate::Take;
+    stack: &mut Stack,
+    call: &Call,
+    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 ctrlc = engine_state.ctrlc.clone();
+    let metadata = input.metadata();
+
+    let mut input_peek = input.into_iter().peekable();
+    if input_peek.peek().is_some() {
+        match input_peek
+            .peek()
+            .ok_or_else(|| {
+                ShellError::GenericError(
+                    "Error in first".into(),
+                    "unable to pick on next value".into(),
+                    Some(call.head),
+                    None,
+                    Vec::new(),
+                )
+            })?
+            .get_type()
+        {
+            Type::Binary => {
+                match &mut input_peek.next() {
+                    Some(v) => match &v {
+                        Value::Binary { val, .. } => {
+                            let bytes = val;
+                            if bytes.len() >= rows_desired {
+                                // We only want to see a certain amount of the binary
+                                // so let's grab those parts
+                                let output_bytes = bytes[0..rows_desired].to_vec();
+                                Ok(Value::Binary {
+                                    val: output_bytes,
+                                    span: head,
+                                }
+                                .into_pipeline_data())
+                            } else {
+                                // if we want more rows that the current chunk size (8192)
+                                // we must gradually get bigger chunks while testing
+                                // if it's within the requested rows_desired size
+                                let mut bigger: Vec<u8> = vec![];
+                                bigger.extend(bytes);
+                                while bigger.len() < rows_desired {
+                                    match input_peek.next() {
+                                        Some(Value::Binary { val, .. }) => bigger.extend(val),
+                                        _ => {
+                                            // We're at the end of our data so let's break out of this loop
+                                            // and set the rows_desired to the size of our data
+                                            rows_desired = bigger.len();
+                                            break;
+                                        }
+                                    }
+                                }
+                                let output_bytes = bigger[0..rows_desired].to_vec();
+                                Ok(Value::Binary {
+                                    val: output_bytes,
+                                    span: head,
+                                }
+                                .into_pipeline_data())
+                            }
+                        }
+
+                        _ => todo!(),
+                    },
+                    None => Ok(input_peek
+                        .into_iter()
+                        .take(rows_desired)
+                        .into_pipeline_data(ctrlc)
+                        .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))
+                }
+            }
+        }
+    } else {
+        Err(ShellError::UnsupportedInput(
+            String::from("Cannot perform into string on empty input"),
+            head,
+        ))
+    }
+}
+#[cfg(test)]
+mod test {
+    use super::*;
     #[test]
     fn test_examples() {
         use crate::test_examples;