Add --skip flag to nth command (#2953)

clippy & rustfmt included
This commit is contained in:
Jakub Žádník 2021-01-20 19:37:30 +02:00 committed by GitHub
parent b435075e09
commit 05e42381df
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -8,6 +8,7 @@ use nu_source::Tagged;
struct NthArgs { struct NthArgs {
row_number: Tagged<u64>, row_number: Tagged<u64>,
rest: Vec<Tagged<u64>>, rest: Vec<Tagged<u64>>,
skip: bool,
} }
pub struct Nth; pub struct Nth;
@ -26,10 +27,11 @@ impl WholeStreamCommand for Nth {
"the number of the row to return", "the number of the row to return",
) )
.rest(SyntaxShape::Any, "Optionally return more rows") .rest(SyntaxShape::Any, "Optionally return more rows")
.switch("skip", "Skip the rows instead of selecting them", Some('s'))
} }
fn usage(&self) -> &str { fn usage(&self) -> &str {
"Return only the selected rows" "Return or skip only the selected rows"
} }
async fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> { async fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
@ -48,6 +50,11 @@ impl WholeStreamCommand for Nth {
example: "echo [first second third] | nth 0 2", example: "echo [first second third] | nth 0 2",
result: Some(vec![Value::from("first"), Value::from("third")]), result: Some(vec![Value::from("first"), Value::from("third")]),
}, },
Example {
description: "Skip the first and third rows",
example: "echo [first second third] | nth --skip 0 2",
result: Some(vec![Value::from("second")]),
},
] ]
} }
} }
@ -57,6 +64,7 @@ async fn nth(args: CommandArgs) -> Result<OutputStream, ShellError> {
NthArgs { NthArgs {
row_number, row_number,
rest: and_rows, rest: and_rows,
skip,
}, },
input, input,
) = args.process().await?; ) = args.process().await?;
@ -67,22 +75,14 @@ async fn nth(args: CommandArgs) -> Result<OutputStream, ShellError> {
.map(|x| x.item) .map(|x| x.item)
.collect::<Vec<u64>>(); .collect::<Vec<u64>>();
let max_row_number = row_numbers
.iter()
.max()
.expect("Internal error: should be > 0 row numbers");
Ok(input Ok(input
.take(*max_row_number as usize + 1)
.enumerate() .enumerate()
.filter_map(move |(idx, item)| { .filter_map(move |(idx, item)| {
futures::future::ready( futures::future::ready(if row_numbers.contains(&(idx as u64)) ^ skip {
if row_numbers.iter().any(|requested| *requested == idx as u64) {
Some(ReturnSuccess::value(item)) Some(ReturnSuccess::value(item))
} else { } else {
None None
}, })
)
}) })
.to_output_stream()) .to_output_stream())
} }