diff --git a/src/uu/shuf/src/shuf.rs b/src/uu/shuf/src/shuf.rs index 6a3e325c7..5950fe4e4 100644 --- a/src/uu/shuf/src/shuf.rs +++ b/src/uu/shuf/src/shuf.rs @@ -75,17 +75,15 @@ 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 headcounts = matches + .values_of(options::HEAD_COUNT) + .unwrap_or_default() + .collect(); + match parse_head_count(&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 +150,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 +298,17 @@ fn parse_range(input_range: &str) -> Result<(usize, usize), String> { } } +fn parse_head_count(headcounts: &Vec<&str>) -> 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..dc3355dff 100644 --- a/tests/by-util/test_shuf.rs +++ b/tests/by-util/test_shuf.rs @@ -196,3 +196,17 @@ 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_seq: Vec<&str> = result + .stdout_str() + .split('\n') + .filter(|x| !x.is_empty()) + .collect(); + assert_eq!(result_seq.len(), 5, "Output should have 5 items"); +}