mirror of
https://github.com/nushell/nushell.git
synced 2024-12-23 23:49:44 +01:00
Make take work like first (#5942)
This commit is contained in:
JT
GitHub
parent
2ac5b0480a
commit
c16d8f0d5f
@ -1,11 +1,9 @@
|
||||
use std::convert::TryInto;
|
||||
|
||||
use nu_engine::CallExt;
|
||||
use nu_protocol::ast::Call;
|
||||
use nu_protocol::engine::{Command, EngineState, Stack};
|
||||
use nu_protocol::{
|
||||
ast::Call,
|
||||
engine::{Command, EngineState, Stack},
|
||||
Category, Example, IntoInterruptiblePipelineData, PipelineData, ShellError, Signature, Span,
|
||||
SyntaxShape, Value,
|
||||
Category, Example, IntoInterruptiblePipelineData, IntoPipelineData, PipelineData, ShellError,
|
||||
Signature, Span, SyntaxShape, Type, Value,
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
@ -17,45 +15,17 @@ impl Command for Take {
|
||||
}
|
||||
|
||||
fn signature(&self) -> Signature {
|
||||
Signature::build(self.name())
|
||||
.optional("n", SyntaxShape::Int, "the number of elements to take")
|
||||
Signature::build("take")
|
||||
.optional(
|
||||
"n",
|
||||
SyntaxShape::Int,
|
||||
"starting from the front, the number of elements to return",
|
||||
)
|
||||
.category(Category::Filters)
|
||||
}
|
||||
|
||||
fn usage(&self) -> &str {
|
||||
"Take the first n elements of the input."
|
||||
}
|
||||
|
||||
fn examples(&self) -> Vec<Example> {
|
||||
vec![
|
||||
Example {
|
||||
description: "Take two elements",
|
||||
example: "echo [[editions]; [2015] [2018] [2021]] | take 2",
|
||||
result: Some(Value::List {
|
||||
vals: vec![
|
||||
Value::Record {
|
||||
cols: vec!["editions".to_owned()],
|
||||
vals: vec![Value::test_int(2015)],
|
||||
span: Span::test_data(),
|
||||
},
|
||||
Value::Record {
|
||||
cols: vec!["editions".to_owned()],
|
||||
vals: vec![Value::test_int(2018)],
|
||||
span: Span::test_data(),
|
||||
},
|
||||
],
|
||||
span: Span::test_data(),
|
||||
}),
|
||||
},
|
||||
Example {
|
||||
description: "Take the first value",
|
||||
example: "echo [2 4 6 8] | take",
|
||||
result: Some(Value::List {
|
||||
vals: vec![Value::test_int(2)],
|
||||
span: Span::test_data(),
|
||||
}),
|
||||
},
|
||||
]
|
||||
"Take only the first n elements."
|
||||
}
|
||||
|
||||
fn run(
|
||||
@ -64,38 +34,134 @@ impl Command for Take {
|
||||
stack: &mut Stack,
|
||||
call: &Call,
|
||||
input: PipelineData,
|
||||
) -> Result<PipelineData, ShellError> {
|
||||
let n: Option<Value> = call.opt(engine_state, stack, 0)?;
|
||||
let metadata = input.metadata();
|
||||
) -> Result<nu_protocol::PipelineData, nu_protocol::ShellError> {
|
||||
first_helper(engine_state, stack, call, input)
|
||||
}
|
||||
|
||||
let n: usize = match n {
|
||||
Some(Value::Int { val, span }) => val.try_into().map_err(|err| {
|
||||
ShellError::UnsupportedInput(
|
||||
format!("Could not convert {} to unsigned integer: {}", val, err),
|
||||
span,
|
||||
)
|
||||
})?,
|
||||
Some(_) => {
|
||||
let span = call.head;
|
||||
return Err(ShellError::TypeMismatch("expected integer".into(), span));
|
||||
}
|
||||
None => 1,
|
||||
};
|
||||
|
||||
let ctrlc = engine_state.ctrlc.clone();
|
||||
|
||||
Ok(input
|
||||
.into_iter()
|
||||
.take(n)
|
||||
.into_pipeline_data(ctrlc)
|
||||
.set_metadata(metadata))
|
||||
fn examples(&self) -> Vec<Example> {
|
||||
vec![
|
||||
Example {
|
||||
description: "Return the first item of a list/table",
|
||||
example: "[1 2 3] | take",
|
||||
result: Some(Value::test_int(1)),
|
||||
},
|
||||
Example {
|
||||
description: "Return the first 2 items of a list/table",
|
||||
example: "[1 2 3] | take 2",
|
||||
result: Some(Value::List {
|
||||
vals: vec![Value::test_int(1), Value::test_int(2)],
|
||||
span: Span::test_data(),
|
||||
}),
|
||||
},
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::Take;
|
||||
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 rows: Option<i64> = call.opt(engine_state, stack, 0)?;
|
||||
let mut rows_desired: usize = match rows {
|
||||
Some(x) => x as usize,
|
||||
None => 1,
|
||||
};
|
||||
|
||||
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)),
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
if rows_desired == 1 {
|
||||
match input_peek.next() {
|
||||
Some(val) => Ok(val.into_pipeline_data()),
|
||||
None => Err(ShellError::AccessBeyondEndOfStream(head)),
|
||||
}
|
||||
} else {
|
||||
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::*;
|
||||
#[test]
|
||||
fn test_examples() {
|
||||
use crate::test_examples;
|
||||
|
||||
Loading…
Reference in New Issue
Block a user