Return errors on unexpected inputs to take and first (#7123)

* Fix `take` behaviour for unexpected input types

* Fix `first` behaviour for unexpected input types

* Fix copy paste mistake
This commit is contained in:
Reilly Wood 2022-11-13 15:15:27 -08:00 committed by GitHub
parent 35f9299fc6
commit 336df6c65e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 177 additions and 168 deletions

View File

@ -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<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 {
// 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<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())
}
}
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<u8> = 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)]

View File

@ -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<nu_protocol::PipelineData, nu_protocol::ShellError> {
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<u8> = 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<Example> {
@ -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<nu_protocol::PipelineData, nu_protocol::ShellError> {
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<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)),
}
}
_ => 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::*;

View File

@ -79,3 +79,16 @@ fn gets_first_row_as_list_when_amount_given() {
assert_eq!(actual.out, "list<int>");
}
#[test]
// covers a situation where `first` used to behave strangely on list<binary> input
fn works_with_binary_list() {
let actual = nu!(
cwd: ".", pipeline(
r#"
([0x[01 11]] | first) == 0x[01 11]
"#
));
assert_eq!(actual.out, "true");
}

View File

@ -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<binary> 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");
}

View File

@ -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<Span> {
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),

View File

@ -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");