diff --git a/crates/nu-command/src/filters/first.rs b/crates/nu-command/src/filters/first.rs index 80430ea1da..2f4ac233ee 100644 --- a/crates/nu-command/src/filters/first.rs +++ b/crates/nu-command/src/filters/first.rs @@ -72,7 +72,7 @@ impl Command for First { }), }, Example { - description: "Return the first 2 items of a bytes", + description: "Return the first 2 bytes of a binary value", example: "0x[01 23 45] | first 2", result: Some(Value::Binary { val: vec![0x01, 0x23], @@ -91,7 +91,12 @@ fn first_helper( ) -> Result { let head = call.head; let rows: Option = call.opt(engine_state, stack, 0)?; - let mut rows_desired: usize = match rows { + // FIXME: for backwards compatibility reasons, if `rows` is not specified we + // return a single element and otherwise we return a single list. We should probably + // remove `rows` so that `first` always returns a single element; getting a list of + // the first N elements is covered by `take` + let return_single_element = rows.is_none(); + let rows_desired: usize = match rows { Some(x) => x as usize, None => 1, }; @@ -99,87 +104,71 @@ fn first_helper( 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()) - } - } + let input_span = input.span(); + let input_not_supported_error = || -> ShellError { + // can't always get a span for input, so try our best and fall back on the span for the `first` call if needed + if let Some(span) = input_span { + ShellError::UnsupportedInput("first does not support this input type".into(), span) + } else { + ShellError::UnsupportedInput( + "first was given an unsupported input type".into(), + call.span(), + ) + } + }; - _ => todo!(), - }, - None => Ok(input_peek - .into_iter() - .take(rows_desired) - .into_pipeline_data(ctrlc) - .set_metadata(metadata)), - } - } - _ => { - if rows_desired == 1 && rows.is_none() { - match input_peek.next() { - Some(val) => Ok(val.into_pipeline_data()), - None => Err(ShellError::AccessBeyondEndOfStream(head)), + match input { + PipelineData::Value(val, _) => match val { + Value::List { vals, .. } => { + if return_single_element { + if vals.is_empty() { + Err(ShellError::AccessEmptyContent(head)) + } else { + Ok(vals[0].clone().into_pipeline_data()) } } else { - Ok(input_peek + Ok(vals .into_iter() .take(rows_desired) .into_pipeline_data(ctrlc) .set_metadata(metadata)) } } + Value::Binary { val, span } => { + let slice: Vec = val.into_iter().take(rows_desired).collect(); + Ok(PipelineData::Value( + Value::Binary { val: slice, span }, + metadata, + )) + } + Value::Range { val, .. } => { + if return_single_element { + Ok(val.from.into_pipeline_data()) + } else { + Ok(val + .into_range_iter(ctrlc.clone())? + .take(rows_desired) + .into_pipeline_data(ctrlc) + .set_metadata(metadata)) + } + } + _ => Err(input_not_supported_error()), + }, + PipelineData::ListStream(mut ls, metadata) => { + if return_single_element { + if let Some(v) = ls.next() { + Ok(v.into_pipeline_data()) + } else { + Err(ShellError::AccessEmptyContent(head)) + } + } else { + Ok(ls + .take(rows_desired) + .into_pipeline_data(ctrlc) + .set_metadata(metadata)) + } } - } else { - Ok(PipelineData::new(head).set_metadata(metadata)) + _ => Err(input_not_supported_error()), } } #[cfg(test)] diff --git a/crates/nu-command/src/filters/take/take_.rs b/crates/nu-command/src/filters/take/take_.rs index 356c42675b..5fdd825127 100644 --- a/crates/nu-command/src/filters/take/take_.rs +++ b/crates/nu-command/src/filters/take/take_.rs @@ -2,8 +2,8 @@ use nu_engine::CallExt; use nu_protocol::ast::Call; use nu_protocol::engine::{Command, EngineState, Stack}; use nu_protocol::{ - Category, Example, IntoInterruptiblePipelineData, IntoPipelineData, PipelineData, ShellError, - Signature, Span, SyntaxShape, Type, Value, + Category, Example, IntoInterruptiblePipelineData, PipelineData, ShellError, Signature, Span, + SyntaxShape, Type, Value, }; #[derive(Clone)] @@ -22,6 +22,8 @@ impl Command for Take { Type::List(Box::new(Type::Any)), Type::List(Box::new(Type::Any)), ), + (Type::Binary, Type::Binary), + (Type::Range, Type::List(Box::new(Type::Number))), ]) .required( "n", @@ -32,7 +34,7 @@ impl Command for Take { } fn usage(&self) -> &str { - "Take only the first n elements." + "Take only the first n elements of a list, or the first n bytes of a binary value." } fn search_terms(&self) -> Vec<&str> { @@ -46,7 +48,51 @@ impl Command for Take { call: &Call, input: PipelineData, ) -> Result { - first_helper(engine_state, stack, call, input) + let rows_desired: usize = call.req(engine_state, stack, 0)?; + + let ctrlc = engine_state.ctrlc.clone(); + let metadata = input.metadata(); + + let input_span = input.span(); + let input_not_supported_error = || -> ShellError { + // can't always get a span for input, so try our best and fall back on the span for the `take` call if needed + if let Some(span) = input_span { + ShellError::UnsupportedInput("take does not support this input type".into(), span) + } else { + ShellError::UnsupportedInput( + "take was given an unsupported input type".into(), + call.span(), + ) + } + }; + + match input { + PipelineData::Value(val, _) => match val { + Value::List { vals, .. } => Ok(vals + .into_iter() + .take(rows_desired) + .into_pipeline_data(ctrlc) + .set_metadata(metadata)), + Value::Binary { val, span } => { + let slice: Vec = val.into_iter().take(rows_desired).collect(); + Ok(PipelineData::Value( + Value::Binary { val: slice, span }, + metadata, + )) + } + Value::Range { val, .. } => Ok(val + .into_range_iter(ctrlc.clone())? + .take(rows_desired) + .into_pipeline_data(ctrlc) + .set_metadata(metadata)), + _ => Err(input_not_supported_error()), + }, + PipelineData::ListStream(ls, metadata) => Ok(ls + .take(rows_desired) + .into_pipeline_data(ctrlc) + .set_metadata(metadata)), + _ => Err(input_not_supported_error()), + } } fn examples(&self) -> Vec { @@ -78,99 +124,26 @@ impl Command for Take { span: Span::test_data(), }), }, + Example { + description: "Return the first 2 bytes of a binary value", + example: "0x[01 23 45] | take 2", + result: Some(Value::Binary { + val: vec![0x01, 0x23], + span: Span::test_data(), + }), + }, + Example { + description: "Return the first 3 elements of a range", + example: "1..10 | take 3", + result: Some(Value::List { + vals: vec![Value::test_int(1), Value::test_int(2), Value::test_int(3)], + span: Span::test_data(), + }), + }, ] } } -fn first_helper( - engine_state: &EngineState, - stack: &mut Stack, - call: &Call, - input: PipelineData, -) -> Result { - let head = call.head; - let mut rows_desired: usize = call.req(engine_state, stack, 0)?; - - 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)), - } - } - _ => 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::*; diff --git a/crates/nu-command/tests/commands/first.rs b/crates/nu-command/tests/commands/first.rs index 8570fe7296..620c206f06 100644 --- a/crates/nu-command/tests/commands/first.rs +++ b/crates/nu-command/tests/commands/first.rs @@ -79,3 +79,16 @@ fn gets_first_row_as_list_when_amount_given() { assert_eq!(actual.out, "list"); } + +#[test] +// covers a situation where `first` used to behave strangely on list input +fn works_with_binary_list() { + let actual = nu!( + cwd: ".", pipeline( + r#" + ([0x[01 11]] | first) == 0x[01 11] + "# + )); + + assert_eq!(actual.out, "true"); +} diff --git a/crates/nu-command/tests/commands/take/rows.rs b/crates/nu-command/tests/commands/take/rows.rs index 3771781d1e..f496acd70c 100644 --- a/crates/nu-command/tests/commands/take/rows.rs +++ b/crates/nu-command/tests/commands/take/rows.rs @@ -41,3 +41,28 @@ fn rows_with_no_arguments_should_lead_to_error() { assert!(actual.err.contains("missing_positional")); }) } + +#[test] +fn fails_on_string() { + let actual = nu!( + cwd: ".", pipeline( + r#" + "foo bar" | take 2 + "# + )); + + assert!(actual.err.contains("unsupported_input")); +} + +#[test] +// covers a situation where `take` used to behave strangely on list input +fn works_with_binary_list() { + let actual = nu!( + cwd: ".", pipeline( + r#" + ([0x[01 11]] | take 1 | get 0) == 0x[01 11] + "# + )); + + assert_eq!(actual.out, "true"); +} diff --git a/crates/nu-protocol/src/pipeline_data.rs b/crates/nu-protocol/src/pipeline_data.rs index 2699bab701..2f266118bb 100644 --- a/crates/nu-protocol/src/pipeline_data.rs +++ b/crates/nu-protocol/src/pipeline_data.rs @@ -89,6 +89,15 @@ impl PipelineData { matches!(self, PipelineData::Value(Value::Nothing { .. }, ..)) } + /// PipelineData doesn't always have a Span, but we can try! + pub fn span(&self) -> Option { + match self { + PipelineData::ListStream(..) => None, + PipelineData::ExternalStream { span, .. } => Some(*span), + PipelineData::Value(v, _) => v.span().ok(), + } + } + pub fn into_value(self, span: Span) -> Value { match self { PipelineData::Value(Value::Nothing { .. }, ..) => Value::nothing(span), diff --git a/tests/plugins/core_inc.rs b/tests/plugins/core_inc.rs index 6bfe79172e..8c38a35977 100644 --- a/tests/plugins/core_inc.rs +++ b/tests/plugins/core_inc.rs @@ -7,7 +7,7 @@ fn chooses_highest_increment_if_given_more_than_one() { let actual = nu_with_plugins!( cwd: "tests/fixtures/formats", plugin: ("nu_plugin_inc"), - "open cargo_sample.toml | first | inc package.version --major --minor | get package.version" + "open cargo_sample.toml | inc package.version --major --minor | get package.version" ); assert_eq!(actual.out, "1.0.0"); @@ -16,7 +16,7 @@ fn chooses_highest_increment_if_given_more_than_one() { cwd: "tests/fixtures/formats", plugin: ("nu_plugin_inc"), // Regardless of order of arguments - "open cargo_sample.toml | first | inc package.version --minor --major | get package.version" + "open cargo_sample.toml | inc package.version --minor --major | get package.version" ); assert_eq!(actual.out, "1.0.0");