shuf: accept multiple occurances of head-count argument

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) {
+        head_count: {
-            Some(count) => match count.parse::<usize>() {
+            let headcounts = matches
+                .values_of(options::HEAD_COUNT)
+                .unwrap_or_default()
+                .collect();
+            match parse_head_count(&headcounts) {
                 Ok(val) => val,
-                Err(_) => {
+                Err(msg) => return Err(USimpleError::new(1, msg)),
-                    return Err(USimpleError::new(
+            }
-                        1,
-                        format!("invalid line count: {}", count.quote()),
-                    ));
-                }
-            },
-            None => std::usize::MAX,
         },
         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<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),

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