Merge pull request #3248 from sylvestre/nproc

nproc: Improve the GNU compat

src/uu/nproc/src/nproc.rs

@@ -32,8 +32,8 @@ const USAGE: &str = "{} [OPTIONS]...";
 pub fn uumain(args: impl uucore::Args) -> UResult<()> {
     let matches = uu_app().get_matches_from(args);
 
-    let mut ignore = match matches.value_of(OPT_IGNORE) {
-        Some(numstr) => match numstr.parse() {
+    let ignore = match matches.value_of(OPT_IGNORE) {
+        Some(numstr) => match numstr.trim().parse() {
             Ok(num) => num,
             Err(e) => {
                 return Err(USimpleError::new(
@@ -45,18 +45,18 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> {
         None => 0,
     };
 
-    if !matches.is_present(OPT_ALL) {
-        // OMP_NUM_THREADS doesn't have an impact on --all
-        ignore += match env::var("OMP_NUM_THREADS") {
-            Ok(threadstr) => threadstr.parse().unwrap_or(0),
-            Err(_) => 0,
-        };
-    }
-
     let mut cores = if matches.is_present(OPT_ALL) {
         num_cpus_all()
     } else {
-        num_cpus::get()
+        // OMP_NUM_THREADS doesn't have an impact on --all
+        match env::var("OMP_NUM_THREADS") {
+            // Uses the OpenMP variable to force the number of threads
+            // If the parsing fails, returns the number of CPU
+            Ok(threadstr) => threadstr.parse().unwrap_or_else(|_| num_cpus::get()),
+            // the variable 'OMP_NUM_THREADS' doesn't exit
+            // fallback to the regular CPU detection
+            Err(_) => num_cpus::get(),
+        }
     };
 
     if cores <= ignore {

tests/by-util/test_nproc.rs

@@ -1,3 +1,4 @@
+// spell-checker:ignore incorrectnumber
 use crate::common::util::*;
 
 #[test]
@@ -15,11 +16,11 @@ fn test_nproc_all_omp() {
 
     let result = TestScenario::new(util_name!())
         .ucmd_keepenv()
-        .env("OMP_NUM_THREADS", "1")
+        .env("OMP_NUM_THREADS", "60")
         .succeeds();
 
     let nproc_omp: u8 = result.stdout_str().trim().parse().unwrap();
-    assert!(nproc - 1 == nproc_omp);
+    assert!(nproc_omp == 60);
 
     let result = TestScenario::new(util_name!())
         .ucmd_keepenv()
@@ -28,20 +29,46 @@ fn test_nproc_all_omp() {
         .succeeds();
     let nproc_omp: u8 = result.stdout_str().trim().parse().unwrap();
     assert!(nproc == nproc_omp);
+
+    // If the parsing fails, returns the number of CPU
+    let result = TestScenario::new(util_name!())
+        .ucmd_keepenv()
+        .env("OMP_NUM_THREADS", "incorrectnumber") // returns the number CPU
+        .succeeds();
+    let nproc_omp: u8 = result.stdout_str().trim().parse().unwrap();
+    assert!(nproc == nproc_omp);
 }
 
 #[test]
 fn test_nproc_ignore() {
     let result = new_ucmd!().succeeds();
-    let nproc: u8 = result.stdout_str().trim().parse().unwrap();
-    if nproc > 1 {
+    let nproc_total: u8 = result.stdout_str().trim().parse().unwrap();
+    if nproc_total > 1 {
         // Ignore all CPU but one
         let result = TestScenario::new(util_name!())
             .ucmd_keepenv()
             .arg("--ignore")
-            .arg((nproc - 1).to_string())
+            .arg((nproc_total - 1).to_string())
             .succeeds();
         let nproc: u8 = result.stdout_str().trim().parse().unwrap();
         assert!(nproc == 1);
+        // Ignore all CPU but one with a string
+        let result = TestScenario::new(util_name!())
+            .ucmd_keepenv()
+            .arg("--ignore= 1")
+            .succeeds();
+        let nproc: u8 = result.stdout_str().trim().parse().unwrap();
+        assert!(nproc_total - 1 == nproc);
     }
 }
+
+#[test]
+fn test_nproc_ignore_all_omp() {
+    let result = TestScenario::new(util_name!())
+        .ucmd_keepenv()
+        .env("OMP_NUM_THREADS", "42")
+        .arg("--ignore=40")
+        .succeeds();
+    let nproc: u8 = result.stdout_str().trim().parse().unwrap();
+    assert!(nproc == 2);
+}