Merge pull request #3329 from DevSabb/shuf-gnu-test

shuf: accept multiple occurances of head-count argument
uutils · Mar 28, 2022 · 746e35c · 746e35c
2 parents d6fd701 + 36ec76e
commit 746e35c
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 11 deletions.
diff --git 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::<usize>() {
+        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<usize, String> {
+    let mut result = std::usize::MAX;
+    for count in headcounts {
+        match count.parse::<usize>() {
+            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<File>),
     RngDefault(rand::rngs::ThreadRng),

diff --git 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");
+}