diff --git a/src/uu/shuf/src/shuf.rs b/src/uu/shuf/src/shuf.rs index 6a3e325c7..0b62ec84a 100644 --- a/src/uu/shuf/src/shuf.rs +++ b/src/uu/shuf/src/shuf.rs @@ -7,7 +7,7 @@ // spell-checker:ignore (ToDO) cmdline evec seps rvec fdata -use clap::{crate_version, Arg, Command}; +use clap::{crate_version, Arg, Command, Values}; use rand::prelude::SliceRandom; use rand::RngCore; use std::fs::File; @@ -75,17 +75,13 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> { }; let options = Options { - head_count: match matches.value_of(options::HEAD_COUNT) { - Some(count) => match count.parse::() { + head_count: { + let mut headcounts: Values<'_> = + matches.values_of(options::HEAD_COUNT).unwrap_or_default(); + match parse_head_count(&mut headcounts) { Ok(val) => val, - Err(_) => { - return Err(USimpleError::new( - 1, - format!("invalid line count: {}", count.quote()), - )); - } - }, - None => std::usize::MAX, + Err(msg) => return Err(USimpleError::new(1, msg)), + } }, output: matches.value_of(options::OUTPUT).map(String::from), random_source: matches.value_of(options::RANDOM_SOURCE).map(String::from), @@ -152,6 +148,7 @@ pub fn uu_app<'a>() -> Command<'a> { .short('n') .long(options::HEAD_COUNT) .takes_value(true) + .multiple_occurrences(true) .value_name("COUNT") .help("output at most COUNT lines"), ) @@ -299,6 +296,17 @@ fn parse_range(input_range: &str) -> Result<(usize, usize), String> { } } +fn parse_head_count(headcounts: &mut Values<'_>) -> Result { + let mut result = std::usize::MAX; + for count in headcounts { + match count.parse::() { + Ok(pv) => result = std::cmp::min(result, pv), + Err(_) => return Err(format!("invalid line count: {}", count.quote())), + } + } + Ok(result) +} + enum WrappedRng { RngFile(rand_read_adapter::ReadRng), RngDefault(rand::rngs::ThreadRng), diff --git a/tests/by-util/test_shuf.rs b/tests/by-util/test_shuf.rs index 86828dc45..682b0dab6 100644 --- a/tests/by-util/test_shuf.rs +++ b/tests/by-util/test_shuf.rs @@ -196,3 +196,19 @@ fn test_shuf_invalid_input_line_count() { .fails() .stderr_contains("invalid line count: 'a'"); } + +#[test] +fn test_shuf_multiple_input_line_count() { + let result = new_ucmd!() + .args(&["-i10-200", "-n", "10", "-n", "5"]) + .succeeds(); + + result.no_stderr(); + + let result_count = result + .stdout_str() + .split('\n') + .filter(|x| !x.is_empty()) + .count(); + assert_eq!(result_count, 5, "Output should have 5 items"); +}