diff --git a/crates/nu-cli/src/commands/drop.rs b/crates/nu-cli/src/commands/drop.rs index 746eb9492d..2d38beafc1 100644 --- a/crates/nu-cli/src/commands/drop.rs +++ b/crates/nu-cli/src/commands/drop.rs @@ -35,20 +35,7 @@ impl WholeStreamCommand for Drop { args: CommandArgs, registry: &CommandRegistry, ) -> Result { - let (DropArgs { rows }, input) = args.process(®istry).await?; - let mut v: Vec<_> = input.into_vec().await; - - let rows_to_drop = if let Some(quantity) = rows { - *quantity as usize - } else { - 1 - }; - - for _ in 0..rows_to_drop { - v.pop(); - } - - Ok(futures::stream::iter(v).to_output_stream()) + drop(args, registry).await } fn examples(&self) -> Vec { @@ -70,6 +57,31 @@ impl WholeStreamCommand for Drop { } } +async fn drop(args: CommandArgs, registry: &CommandRegistry) -> Result { + let (DropArgs { rows }, input) = args.process(®istry).await?; + let v: Vec<_> = input.into_vec().await; + + let rows_to_drop = if let Some(quantity) = rows { + *quantity as usize + } else { + 1 + }; + + Ok(if rows_to_drop == 0 { + futures::stream::iter(v).to_output_stream() + } else { + let k = if v.len() < rows_to_drop { + 0 + } else { + v.len() - rows_to_drop + }; + + let iter = v.into_iter().take(k); + + futures::stream::iter(iter).to_output_stream() + }) +} + #[cfg(test)] mod tests { use super::Drop; diff --git a/crates/nu-cli/tests/commands/drop.rs b/crates/nu-cli/tests/commands/drop.rs index 432f94f68f..d33afc6776 100644 --- a/crates/nu-cli/tests/commands/drop.rs +++ b/crates/nu-cli/tests/commands/drop.rs @@ -1,4 +1,4 @@ -use nu_test_support::nu; +use nu_test_support::{nu, pipeline}; #[test] fn drop_rows() { @@ -9,3 +9,15 @@ fn drop_rows() { assert_eq!(actual.out, "3"); } + +#[test] +fn drop_more_rows_than_table_has() { + let actual = nu!( + cwd: ".", pipeline( + r#" + date | drop 50 | count + "# + )); + + assert_eq!(actual.out, "0"); +}