nushell/crates/nu-cli/src/commands/nth.rs
Jonathan Turner ac578b8491
Multiline scripts part 2 (#2795)
* Begin allowing comments and multiline scripts.

* clippy

* Finish moving to groups. Test pass

* Keep going

* WIP

* WIP

* BROKEN WIP

* WIP

* WIP

* Fix more tests

* WIP: alias starts working

* Broken WIP

* Broken WIP

* Variables begin to work

* captures start working

* A little better but needs fixed scope

* Shorthand env setting

* Update main merge

* Broken WIP

* WIP

* custom command parsing

* Custom commands start working

* Fix coloring and parsing of block

* Almost there

* Add some tests

* Add more param types

* Bump version

* Fix benchmark

* Fix stuff
2020-12-18 20:53:49 +13:00

102 lines
2.5 KiB
Rust

use crate::commands::WholeStreamCommand;
use crate::prelude::*;
use nu_errors::ShellError;
use nu_protocol::{ReturnSuccess, Signature, SyntaxShape, Value};
use nu_source::Tagged;
#[derive(Deserialize)]
struct NthArgs {
row_number: Tagged<u64>,
rest: Vec<Tagged<u64>>,
}
pub struct Nth;
#[async_trait]
impl WholeStreamCommand for Nth {
fn name(&self) -> &str {
"nth"
}
fn signature(&self) -> Signature {
Signature::build("nth")
.required(
"row number",
SyntaxShape::Int,
"the number of the row to return",
)
.rest(SyntaxShape::Any, "Optionally return more rows")
}
fn usage(&self) -> &str {
"Return only the selected rows"
}
async fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
nth(args).await
}
fn examples(&self) -> Vec<Example> {
vec![
Example {
description: "Get the second row",
example: "echo [first second third] | nth 1",
result: Some(vec![Value::from("second")]),
},
Example {
description: "Get the first and third rows",
example: "echo [first second third] | nth 0 2",
result: Some(vec![Value::from("first"), Value::from("third")]),
},
]
}
}
async fn nth(args: CommandArgs) -> Result<OutputStream, ShellError> {
let (
NthArgs {
row_number,
rest: and_rows,
},
input,
) = args.process().await?;
let row_numbers = vec![vec![row_number], and_rows]
.into_iter()
.flatten()
.map(|x| x.item)
.collect::<Vec<u64>>();
let max_row_number = row_numbers
.iter()
.max()
.expect("Internal error: should be > 0 row numbers");
Ok(input
.take(*max_row_number as usize + 1)
.enumerate()
.filter_map(move |(idx, item)| {
futures::future::ready(
if row_numbers.iter().any(|requested| *requested == idx as u64) {
Some(ReturnSuccess::value(item))
} else {
None
},
)
})
.to_output_stream())
}
#[cfg(test)]
mod tests {
use super::Nth;
use super::ShellError;
#[test]
fn examples_work_as_expected() -> Result<(), ShellError> {
use crate::examples::test as test_examples;
Ok(test_examples(Nth {})?)
}
}