Fix drop nth bug (#5312)

* Fix drop nth bug on ranges. Should fix & close #5260

* Fix drop nth bug on ranges. Should fix & close #5260

* Add support for ranges

* Working version of drop nth, but the issue is that we unwrap the value which is problematic for Streams. Should convert to the way @stormasm was doing it before and implement the range check

* Fix fmt issue

* Drop nth now works for Lists, Records, and Ranges. We need support for ListStreams and for ExternalStreams

* Keep consistent naming

* Fix fmt issue

* Support ListStreams for drop nth

* Use DropNthIterator instead

* Found a more elegant way to deal with the check for no upper bound input

* Add extra checks for negative inputs or to < from for ranges

Co-authored-by: Stefan Stanciulescu <test@test.com>
This commit is contained in:
Stefan Stanciulescu 2022-06-13 20:49:59 -07:00 committed by GitHub
parent de554f8e5f
commit dc1248a454
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -77,6 +77,22 @@ impl Command for DropNth {
span: Span::test_data(), span: Span::test_data(),
}), }),
}, },
Example {
example: "[0,1,2,3,4,5] | drop nth 1..",
description: "Drop all rows except first row",
result: Some(Value::List {
vals: vec![Value::test_int(0)],
span: Span::test_data(),
}),
},
Example {
example: "[0,1,2,3,4,5] | drop nth 3..",
description: "Drop rows 3,4,5",
result: Some(Value::List {
vals: vec![Value::test_int(0), Value::test_int(1), Value::test_int(2)],
span: Span::test_data(),
}),
},
] ]
} }
@ -87,25 +103,47 @@ impl Command for DropNth {
call: &Call, call: &Call,
input: PipelineData, input: PipelineData,
) -> Result<PipelineData, ShellError> { ) -> Result<PipelineData, ShellError> {
// let mut rows: Vec<usize> = call.rest(engine_state, stack, 0)?;
// rows.sort_unstable();
// let pipeline_iter: PipelineIterator = input.into_iter();
let number_or_range = extract_int_or_range(engine_state, stack, call)?; let number_or_range = extract_int_or_range(engine_state, stack, call)?;
let mut lower_bound = None;
let rows = match number_or_range { let rows = match number_or_range {
Either::Left(row_number) => { Either::Left(row_number) => {
let and_rows: Vec<Spanned<i64>> = call.rest(engine_state, stack, 1)?; let and_rows: Vec<Spanned<i64>> = call.rest(engine_state, stack, 1)?;
let mut rows: Vec<_> = and_rows.into_iter().map(|x| x.item as usize).collect(); let mut rows: Vec<_> = and_rows.into_iter().map(|x| x.item as usize).collect();
rows.push(row_number as usize); rows.push(row_number as usize);
rows.sort_unstable(); rows.sort_unstable();
rows rows
} }
Either::Right(row_range) => { Either::Right(row_range) => {
let from = row_range.from.as_integer()? as usize; let from = row_range.from.as_integer()?; // as usize;
let to = row_range.to.as_integer()? as usize; let to = row_range.to.as_integer()?; // as usize;
if matches!(row_range.inclusion, RangeInclusion::Inclusive) { // check for negative range inputs, e.g., (2..-5)
if from.is_negative() || to.is_negative() {
let span: Spanned<Range> = call.req(engine_state, stack, 0)?;
return Err(ShellError::UnsupportedInput(
"Drop nth accepts only positive integers".to_string(),
span.span,
));
}
// check if the upper bound is smaller than the lower bound, e.g., do not accept 4..2
if to < from {
let span: Spanned<Range> = call.req(engine_state, stack, 0)?;
return Err(ShellError::UnsupportedInput(
"The upper bound needs to be equal or larger to the lower bound"
.to_string(),
span.span,
));
}
// check for equality to isize::MAX because for some reason,
// the parser returns isize::MAX when we provide a range without upper bound (e.g., 5.. )
let to = to as usize;
let from = from as usize;
if to > 0 && to as isize == isize::MAX {
lower_bound = Some(from);
vec![from]
} else if matches!(row_range.inclusion, RangeInclusion::Inclusive) {
(from..=to).collect() (from..=to).collect()
} else { } else {
(from..to).collect() (from..to).collect()
@ -113,6 +151,13 @@ impl Command for DropNth {
} }
}; };
if let Some(lower_bound) = lower_bound {
Ok(input
.into_iter()
.take(lower_bound)
.collect::<Vec<_>>()
.into_pipeline_data(engine_state.ctrlc.clone()))
} else {
Ok(DropNthIterator { Ok(DropNthIterator {
input: input.into_iter(), input: input.into_iter(),
rows, rows,
@ -121,6 +166,7 @@ impl Command for DropNth {
.into_pipeline_data(engine_state.ctrlc.clone())) .into_pipeline_data(engine_state.ctrlc.clone()))
} }
} }
}
fn extract_int_or_range( fn extract_int_or_range(
engine_state: &EngineState, engine_state: &EngineState,