split: accept leading '+' in the number of chunks

src/uu/split/src/strategy.rs

@@ -117,13 +117,17 @@ impl NumberType {
     /// or if `K` is greater than `N`
     /// then this function returns [`NumberTypeError`].
     fn from(s: &str) -> Result<Self, NumberTypeError> {
+        /// GNU accepts a leading '+' on every number, so strip it first.
+        fn parse_number(value: &str) -> Result<u64, ParseSizeError> {
+            parse_size_u64(value.strip_prefix('+').unwrap_or(value))
+        }
         fn is_invalid_chunk(chunk_number: u64, num_chunks: u64) -> bool {
             chunk_number > num_chunks || chunk_number == 0
         }
         let mut parts = s.splitn(4, '/');
         match (parts.next(), parts.next(), parts.next(), parts.next()) {
             (Some(n_str), None, None, None) => {
-                let num_chunks = parse_size_u64(n_str)
+                let num_chunks = parse_number(n_str)
                     .map_err(|_| NumberTypeError::NumberOfChunks(n_str.to_string()))?;
                 if num_chunks > 0 {
                     Ok(Self::Bytes(num_chunks))
@@ -134,45 +138,45 @@ impl NumberType {
             (Some(k_str), Some(n_str), None, None)
                 if !k_str.starts_with('l') && !k_str.starts_with('r') =>
             {
-                let num_chunks = parse_size_u64(n_str)
+                let num_chunks = parse_number(n_str)
                     .map_err(|_| NumberTypeError::NumberOfChunks(n_str.to_string()))?;
-                let chunk_number = parse_size_u64(k_str)
+                let chunk_number = parse_number(k_str)
                     .map_err(|_| NumberTypeError::ChunkNumber(k_str.to_string()))?;
                 if is_invalid_chunk(chunk_number, num_chunks) {
                     return Err(NumberTypeError::ChunkNumber(k_str.to_string()));
                 }
                 Ok(Self::KthBytes(chunk_number, num_chunks))
             }
             (Some("l"), Some(n_str), None, None) => {
-                let num_chunks = parse_size_u64(n_str)
+                let num_chunks = parse_number(n_str)
                     .map_err(|_| NumberTypeError::NumberOfChunks(n_str.to_string()))?;
                 if num_chunks == 0 {
                     return Err(NumberTypeError::NumberOfChunks(n_str.to_string()));
                 }
                 Ok(Self::Lines(num_chunks))
             }
             (Some("l"), Some(k_str), Some(n_str), None) => {
-                let num_chunks = parse_size_u64(n_str)
+                let num_chunks = parse_number(n_str)
                     .map_err(|_| NumberTypeError::NumberOfChunks(n_str.to_string()))?;
-                let chunk_number = parse_size_u64(k_str)
+                let chunk_number = parse_number(k_str)
                     .map_err(|_| NumberTypeError::ChunkNumber(k_str.to_string()))?;
                 if is_invalid_chunk(chunk_number, num_chunks) {
                     return Err(NumberTypeError::ChunkNumber(k_str.to_string()));
                 }
                 Ok(Self::KthLines(chunk_number, num_chunks))
             }
             (Some("r"), Some(n_str), None, None) => {
-                let num_chunks = parse_size_u64(n_str)
+                let num_chunks = parse_number(n_str)
                     .map_err(|_| NumberTypeError::NumberOfChunks(n_str.to_string()))?;
                 if num_chunks == 0 {
                     return Err(NumberTypeError::NumberOfChunks(n_str.to_string()));
                 }
                 Ok(Self::RoundRobin(num_chunks))
             }
             (Some("r"), Some(k_str), Some(n_str), None) => {
-                let num_chunks = parse_size_u64(n_str)
+                let num_chunks = parse_number(n_str)
                     .map_err(|_| NumberTypeError::NumberOfChunks(n_str.to_string()))?;
-                let chunk_number = parse_size_u64(k_str)
+                let chunk_number = parse_number(k_str)
                     .map_err(|_| NumberTypeError::ChunkNumber(k_str.to_string()))?;
                 if is_invalid_chunk(chunk_number, num_chunks) {
                     return Err(NumberTypeError::ChunkNumber(k_str.to_string()));

tests/by-util/test_split.rs

@@ -9,6 +9,7 @@ use rand::{RngExt as _, SeedableRng, rng};
 use regex::Regex;
 #[cfg(target_os = "linux")]
 use rlimit::Resource;
+use rstest::rstest;
 #[cfg(target_os = "linux")]
 use std::os::unix::ffi::OsStringExt;
 use std::path::Path;
@@ -2369,3 +2370,42 @@ fn test_obsolete_lines_not_read_after_double_dash() {
         .fails()
         .stderr_contains("cannot open '-1' for reading");
 }
+
+#[rstest]
+#[case::bytes("+2", &[("xaa", "a\nb\n"), ("xab", "c\nd\n")], "")]
+#[case::lines("l/+2", &[("xaa", "a\nb\n"), ("xab", "c\nd\n")], "")]
+#[case::round_robin("r/+2", &[("xaa", "a\nc\n"), ("xab", "b\nd\n")], "")]
+#[case::bytes_first_half("+1/2", &[], "a\nb\n")]
+#[case::lines_first_half("l/+1/2", &[], "a\nb\n")]
+#[case::round_robin_first_half("r/+1/2", &[], "a\nc\n")]
+fn test_number_with_plus_prefix(
+    #[case] number: &str,
+    #[case] expected_files: &[(&str, &str)],
+    #[case] expected_stdout: &str,
+) {
+    let (at, mut ucmd) = at_and_ucmd!();
+    at.write("a.txt", "a\nb\nc\nd\n");
+
+    let stdout = ucmd
+        .args(&["-n", number, "a.txt"])
+        .succeeds()
+        .stdout_move_str();
+    assert_eq!(stdout, expected_stdout, "-n {number}");
+
+    let mut files: Vec<(String, String)> = fs::read_dir(at.as_string())
+        .unwrap()
+        .map(|entry| entry.unwrap().file_name())
+        .filter(|name| name.to_string_lossy().starts_with('x'))
+        .map(|name| {
+            let contents = at.read(at.plus(&name));
+            (name.to_string_lossy().into_owned(), contents)
+        })
+        .collect();
+    files.sort();
+
+    let expected_files: Vec<(String, String)> = expected_files
+        .iter()
+        .map(|(name, contents)| (name.to_string(), contents.to_string()))
+        .collect();
+    assert_eq!(files, expected_files, "-n {number}");
+}