mirror of
https://github.com/nushell/nushell.git
synced 2025-07-01 07:00:37 +02:00
drop nth command supports spreadable arguments (#15897)
## ✅ Improve `drop nth` command to support spreadable arguments ### Summary This PR updates the `drop nth` command to support **spreadable arguments** in a way consistent with other commands like `which`, enabling: ```nu [1 2 3 4 5] | drop nth 0 2 4 ``` ### What's Changed * **Previously**: only a single index or a single range was accepted as the first argument, with rest arguments ignored for ranges. * **Now**: the command accepts any combination of: * Integers: to drop individual rows * Ranges: to drop slices of rows * Unbounded ranges: like `3..`, to drop from index onward Example: ```nu [one two three four five six] | drop nth 0 2 4..5 # drops "one", "three", "five", and "six" ``` ### Test Manual Test:  ### Notes As per feedback: * We **only collect the list of indices** to drop, not the input stream. * Unbounded ranges are handled by terminating the stream early. Let me know if you'd like further changes --------- Co-authored-by: Kumar Ujjawal <kumar.ujjawal@greenpista.com> Co-authored-by: Kumar Ujjawal <kumarujjawal@Kumars-MacBook-Air.local>
This commit is contained in:
@ -1,6 +1,6 @@
|
|||||||
use itertools::Either;
|
|
||||||
use nu_engine::command_prelude::*;
|
use nu_engine::command_prelude::*;
|
||||||
use nu_protocol::{PipelineIterator, Range};
|
use nu_protocol::{PipelineIterator, Range};
|
||||||
|
use std::collections::VecDeque;
|
||||||
use std::ops::Bound;
|
use std::ops::Bound;
|
||||||
|
|
||||||
#[derive(Clone)]
|
#[derive(Clone)]
|
||||||
@ -18,13 +18,11 @@ impl Command for DropNth {
|
|||||||
(Type::list(Type::Any), Type::list(Type::Any)),
|
(Type::list(Type::Any), Type::list(Type::Any)),
|
||||||
])
|
])
|
||||||
.allow_variants_without_examples(true)
|
.allow_variants_without_examples(true)
|
||||||
.required(
|
.rest(
|
||||||
"row number or row range",
|
"rest",
|
||||||
// FIXME: we can make this accept either Int or Range when we can compose SyntaxShapes
|
|
||||||
SyntaxShape::Any,
|
SyntaxShape::Any,
|
||||||
"The number of the row to drop or a range to drop consecutive rows.",
|
"The row numbers or ranges to drop.",
|
||||||
)
|
)
|
||||||
.rest("rest", SyntaxShape::Any, "The number of the row to drop.")
|
|
||||||
.category(Category::Filters)
|
.category(Category::Filters)
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -103,110 +101,125 @@ impl Command for DropNth {
|
|||||||
) -> Result<PipelineData, ShellError> {
|
) -> Result<PipelineData, ShellError> {
|
||||||
let head = call.head;
|
let head = call.head;
|
||||||
let metadata = input.metadata();
|
let metadata = input.metadata();
|
||||||
let number_or_range = extract_int_or_range(engine_state, stack, call)?;
|
|
||||||
|
|
||||||
let rows = match number_or_range.item {
|
let args: Vec<Value> = call.rest(engine_state, stack, 0)?;
|
||||||
Either::Left(row_number) => {
|
if args.is_empty() {
|
||||||
let and_rows: Vec<Spanned<i64>> = call.rest(engine_state, stack, 1)?;
|
return Ok(input);
|
||||||
let mut rows: Vec<_> = and_rows.into_iter().map(|x| x.item as usize).collect();
|
}
|
||||||
rows.push(row_number as usize);
|
|
||||||
rows.sort_unstable();
|
|
||||||
rows
|
|
||||||
}
|
|
||||||
Either::Right(Range::FloatRange(_)) => {
|
|
||||||
return Err(ShellError::UnsupportedInput {
|
|
||||||
msg: "float range".into(),
|
|
||||||
input: "value originates from here".into(),
|
|
||||||
msg_span: head,
|
|
||||||
input_span: number_or_range.span,
|
|
||||||
});
|
|
||||||
}
|
|
||||||
Either::Right(Range::IntRange(range)) => {
|
|
||||||
// check for negative range inputs, e.g., (2..-5)
|
|
||||||
let end_negative = match range.end() {
|
|
||||||
Bound::Included(end) | Bound::Excluded(end) => end < 0,
|
|
||||||
Bound::Unbounded => false,
|
|
||||||
};
|
|
||||||
if range.start().is_negative() || end_negative {
|
|
||||||
return Err(ShellError::UnsupportedInput {
|
|
||||||
msg: "drop nth accepts only positive ints".into(),
|
|
||||||
input: "value originates from here".into(),
|
|
||||||
msg_span: head,
|
|
||||||
input_span: number_or_range.span,
|
|
||||||
});
|
|
||||||
}
|
|
||||||
// check if the upper bound is smaller than the lower bound, e.g., do not accept 4..2
|
|
||||||
if range.step() < 0 {
|
|
||||||
return Err(ShellError::UnsupportedInput {
|
|
||||||
msg: "The upper bound needs to be equal or larger to the lower bound"
|
|
||||||
.into(),
|
|
||||||
input: "value originates from here".into(),
|
|
||||||
msg_span: head,
|
|
||||||
input_span: number_or_range.span,
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
let start = range.start() as usize;
|
let (rows_to_drop, min_unbounded_start) = get_rows_to_drop(&args, head)?;
|
||||||
|
|
||||||
let end = match range.end() {
|
let input = if let Some(cutoff) = min_unbounded_start {
|
||||||
Bound::Included(end) => end as usize,
|
input
|
||||||
Bound::Excluded(end) => (end - 1) as usize,
|
.into_iter()
|
||||||
Bound::Unbounded => {
|
.take(cutoff)
|
||||||
return Ok(input
|
.into_pipeline_data_with_metadata(
|
||||||
.into_iter()
|
head,
|
||||||
.take(start)
|
engine_state.signals().clone(),
|
||||||
.into_pipeline_data_with_metadata(
|
metadata.clone(),
|
||||||
head,
|
)
|
||||||
engine_state.signals().clone(),
|
} else {
|
||||||
metadata,
|
input
|
||||||
));
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let end = if let PipelineData::Value(Value::List { vals, .. }, _) = &input {
|
|
||||||
end.min(vals.len() - 1)
|
|
||||||
} else {
|
|
||||||
end
|
|
||||||
};
|
|
||||||
|
|
||||||
(start..=end).collect()
|
|
||||||
}
|
|
||||||
};
|
};
|
||||||
|
|
||||||
Ok(DropNthIterator {
|
Ok(DropNthIterator {
|
||||||
input: input.into_iter(),
|
input: input.into_iter(),
|
||||||
rows,
|
rows: rows_to_drop,
|
||||||
current: 0,
|
current: 0,
|
||||||
}
|
}
|
||||||
.into_pipeline_data_with_metadata(head, engine_state.signals().clone(), metadata))
|
.into_pipeline_data_with_metadata(head, engine_state.signals().clone(), metadata))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn extract_int_or_range(
|
fn get_rows_to_drop(
|
||||||
engine_state: &EngineState,
|
args: &[Value],
|
||||||
stack: &mut Stack,
|
head: Span,
|
||||||
call: &Call,
|
) -> Result<(VecDeque<usize>, Option<usize>), ShellError> {
|
||||||
) -> Result<Spanned<Either<i64, Range>>, ShellError> {
|
let mut rows_to_drop = Vec::new();
|
||||||
let value: Value = call.req(engine_state, stack, 0)?;
|
let mut min_unbounded_start: Option<usize> = None;
|
||||||
|
|
||||||
let int_opt = value.as_int().map(Either::Left).ok();
|
for value in args {
|
||||||
let range_opt = value.as_range().map(Either::Right).ok();
|
if let Ok(i) = value.as_int() {
|
||||||
|
if i < 0 {
|
||||||
|
return Err(ShellError::UnsupportedInput {
|
||||||
|
msg: "drop nth accepts only positive ints".into(),
|
||||||
|
input: "value originates from here".into(),
|
||||||
|
msg_span: head,
|
||||||
|
input_span: value.span(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
rows_to_drop.push(i as usize);
|
||||||
|
} else if let Ok(range) = value.as_range() {
|
||||||
|
match range {
|
||||||
|
Range::IntRange(range) => {
|
||||||
|
let start = range.start();
|
||||||
|
if start < 0 {
|
||||||
|
return Err(ShellError::UnsupportedInput {
|
||||||
|
msg: "drop nth accepts only positive ints".into(),
|
||||||
|
input: "value originates from here".into(),
|
||||||
|
msg_span: head,
|
||||||
|
input_span: value.span(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
int_opt
|
match range.end() {
|
||||||
.or(range_opt)
|
Bound::Included(end) => {
|
||||||
.ok_or_else(|| ShellError::TypeMismatch {
|
if end < start {
|
||||||
err_message: "int or range".into(),
|
return Err(ShellError::UnsupportedInput {
|
||||||
span: value.span(),
|
msg: "The upper bound must be greater than or equal to the lower bound".into(),
|
||||||
})
|
input: "value originates from here".into(),
|
||||||
.map(|either| Spanned {
|
msg_span: head,
|
||||||
item: either,
|
input_span: value.span(),
|
||||||
span: value.span(),
|
});
|
||||||
})
|
}
|
||||||
|
rows_to_drop.extend((start as usize)..=(end as usize));
|
||||||
|
}
|
||||||
|
Bound::Excluded(end) => {
|
||||||
|
if end <= start {
|
||||||
|
return Err(ShellError::UnsupportedInput {
|
||||||
|
msg: "The upper bound must be greater than the lower bound"
|
||||||
|
.into(),
|
||||||
|
input: "value originates from here".into(),
|
||||||
|
msg_span: head,
|
||||||
|
input_span: value.span(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
rows_to_drop.extend((start as usize)..(end as usize));
|
||||||
|
}
|
||||||
|
Bound::Unbounded => {
|
||||||
|
let start_usize = start as usize;
|
||||||
|
min_unbounded_start = Some(
|
||||||
|
min_unbounded_start.map_or(start_usize, |s| s.min(start_usize)),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Range::FloatRange(_) => {
|
||||||
|
return Err(ShellError::UnsupportedInput {
|
||||||
|
msg: "float range not supported".into(),
|
||||||
|
input: "value originates from here".into(),
|
||||||
|
msg_span: head,
|
||||||
|
input_span: value.span(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
return Err(ShellError::TypeMismatch {
|
||||||
|
err_message: "Expected int or range".into(),
|
||||||
|
span: value.span(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
rows_to_drop.sort_unstable();
|
||||||
|
rows_to_drop.dedup();
|
||||||
|
|
||||||
|
Ok((VecDeque::from(rows_to_drop), min_unbounded_start))
|
||||||
}
|
}
|
||||||
|
|
||||||
struct DropNthIterator {
|
struct DropNthIterator {
|
||||||
input: PipelineIterator,
|
input: PipelineIterator,
|
||||||
rows: Vec<usize>,
|
rows: VecDeque<usize>,
|
||||||
current: usize,
|
current: usize,
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -215,9 +228,9 @@ impl Iterator for DropNthIterator {
|
|||||||
|
|
||||||
fn next(&mut self) -> Option<Self::Item> {
|
fn next(&mut self) -> Option<Self::Item> {
|
||||||
loop {
|
loop {
|
||||||
if let Some(row) = self.rows.first() {
|
if let Some(row) = self.rows.front() {
|
||||||
if self.current == *row {
|
if self.current == *row {
|
||||||
self.rows.remove(0);
|
self.rows.pop_front();
|
||||||
self.current += 1;
|
self.current += 1;
|
||||||
let _ = self.input.next();
|
let _ = self.input.next();
|
||||||
continue;
|
continue;
|
||||||
|
@ -135,3 +135,75 @@ fn more_columns_than_record_has() {
|
|||||||
|
|
||||||
assert_eq!(actual.out, "{}");
|
assert_eq!(actual.out, "{}");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn drop_single_index() {
|
||||||
|
let actual = nu!("echo 10..15 | drop nth 2 | to json --raw");
|
||||||
|
assert_eq!(actual.out, "[10,11,13,14,15]");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn drop_multiple_indices() {
|
||||||
|
let actual = nu!("echo 0..10 | drop nth 1 3 | to json --raw");
|
||||||
|
assert_eq!(actual.out, "[0,2,4,5,6,7,8,9,10]");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn drop_inclusive_range() {
|
||||||
|
let actual = nu!("echo 10..15 | drop nth (2..4) | to json --raw");
|
||||||
|
assert_eq!(actual.out, "[10,11,15]");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn drop_exclusive_range() {
|
||||||
|
let actual = nu!("echo 10..15 | drop nth (2..<4) | to json --raw");
|
||||||
|
assert_eq!(actual.out, "[10,11,14,15]");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn drop_unbounded_range() {
|
||||||
|
let actual = nu!("echo 0..5 | drop nth 3.. | to json --raw");
|
||||||
|
assert_eq!(actual.out, "[0,1,2]");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn drop_multiple_ranges_including_unbounded() {
|
||||||
|
let actual = nu!(pipeline(
|
||||||
|
r#"
|
||||||
|
0..30
|
||||||
|
| drop nth 0..10 20..
|
||||||
|
| to json --raw
|
||||||
|
"#
|
||||||
|
));
|
||||||
|
|
||||||
|
assert_eq!(actual.out, "[11,12,13,14,15,16,17,18,19]");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn drop_combination_of_unbounded_range_and_single_index() {
|
||||||
|
let actual = nu!(pipeline(
|
||||||
|
r#"
|
||||||
|
echo 0..15
|
||||||
|
| drop nth 10.. 5
|
||||||
|
| to json --raw
|
||||||
|
"#
|
||||||
|
));
|
||||||
|
|
||||||
|
assert_eq!(actual.out, "[0,1,2,3,4,6,7,8,9]");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn drop_combination_of_two_unbounded_ranges() {
|
||||||
|
let actual = nu!(pipeline(
|
||||||
|
r#"
|
||||||
|
echo 0..150
|
||||||
|
| drop nth 0..100 999..
|
||||||
|
| to json --raw
|
||||||
|
"#
|
||||||
|
));
|
||||||
|
|
||||||
|
let expected: Vec<u32> = (101..=150).collect();
|
||||||
|
let expected_json = serde_json::to_string(&expected).unwrap();
|
||||||
|
|
||||||
|
assert_eq!(actual.out, expected_json);
|
||||||
|
}
|
||||||
|
Reference in New Issue
Block a user