From c16d8f0d5fb3f22aa596541d4a4dcf6a3e4f9fa5 Mon Sep 17 00:00:00 2001 From: JT <547158+jntrnr@users.noreply.github.com> Date: Mon, 4 Jul 2022 08:03:35 +1200 Subject: [PATCH] Make `take` work like `first` (#5942) --- crates/nu-command/src/filters/take/take_.rs | 202 +++++++++++++------- 1 file changed, 134 insertions(+), 68 deletions(-) diff --git a/crates/nu-command/src/filters/take/take_.rs b/crates/nu-command/src/filters/take/take_.rs index 64f76e53e2..d887f473b4 100644 --- a/crates/nu-command/src/filters/take/take_.rs +++ b/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, - engine::{Command, EngineState, Stack}, - Category, Example, IntoInterruptiblePipelineData, PipelineData, ShellError, Signature, Span, - SyntaxShape, Value, + Category, Example, IntoInterruptiblePipelineData, IntoPipelineData, PipelineData, ShellError, + Signature, Span, SyntaxShape, Type, Value, }; #[derive(Clone)] @@ -17,45 +15,17 @@ impl Command for Take { } fn signature(&self) -> Signature { - Signature::build(self.name()) - .optional("n", SyntaxShape::Int, "the number of elements to take") + Signature::build("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." - } - - fn examples(&self) -> Vec { - 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(), - }), - }, - ] + "Take only the first n elements." } fn run( @@ -64,38 +34,134 @@ impl Command for Take { stack: &mut Stack, call: &Call, input: PipelineData, - ) -> Result { - let n: Option = call.opt(engine_state, stack, 0)?; - let metadata = input.metadata(); + ) -> Result { + first_helper(engine_state, stack, call, input) + } - let n: usize = match n { - Some(Value::Int { val, span }) => val.try_into().map_err(|err| { - ShellError::UnsupportedInput( - format!("Could not convert {} to unsigned integer: {}", val, err), - span, - ) - })?, - Some(_) => { - let span = call.head; - return Err(ShellError::TypeMismatch("expected integer".into(), span)); - } - None => 1, - }; - - let ctrlc = engine_state.ctrlc.clone(); - - Ok(input - .into_iter() - .take(n) - .into_pipeline_data(ctrlc) - .set_metadata(metadata)) + fn examples(&self) -> Vec { + vec![ + Example { + description: "Return the first item of a list/table", + example: "[1 2 3] | take", + result: Some(Value::test_int(1)), + }, + Example { + description: "Return the first 2 items of a list/table", + example: "[1 2 3] | take 2", + result: Some(Value::List { + vals: vec![Value::test_int(1), Value::test_int(2)], + span: Span::test_data(), + }), + }, + ] } } -#[cfg(test)] -mod tests { - use crate::Take; +fn first_helper( + engine_state: &EngineState, + stack: &mut Stack, + call: &Call, + input: PipelineData, +) -> Result { + let head = call.head; + let rows: Option = 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 = 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;