nproc: support --all and --ignore and respect OMP_NUM_THREADS

src/nproc/nproc.rs

@@ -9,6 +9,8 @@
  * file that was distributed with this source code.
  */
 
+#![feature(macro_rules)]
+
 extern crate getopts;
 
 use std::os;
@@ -16,15 +18,23 @@ use std::os;
 static NAME : &'static str = "nproc";
 static VERSION : &'static str = "0.0.0";
 
+#[path = "../common/util.rs"]
+mod util;
+
 pub fn uumain(args: Vec<String>) -> int {
     let opts = [
+        getopts::optflag("", "all", "print the number of cores available to the system"),
+        getopts::optopt("", "ignore", "ignore up to N cores", "N"),
         getopts::optflag("h", "help", "display this help and exit"),
         getopts::optflag("V", "version", "output version information and exit"),
     ];
 
     let matches = match getopts::getopts(args.tail(), opts) {
         Ok(m) => m,
-        Err(err) => fail!("{}", err),
+        Err(err) => {
+            show_error!("{}", err);
+            return 1;
+        }
     };
 
     if matches.opt_present("version") {
@@ -38,14 +48,38 @@ pub fn uumain(args: Vec<String>) -> int {
         println!("Usage:");
         println!("  {} [OPTIONS] NAME...", NAME);
         println!("");
-        print!("{}", getopts::usage("Print the number of cores available.", opts.as_slice()).as_slice());
-        if matches.free.is_empty() {
-            return 1;
-        }
+        print!("{}", getopts::usage("Print the number of cores available to the current process.", opts));
         return 0;
     }
 
-    println!("{}", os::num_cpus());
+    let mut ignore = match matches.opt_str("ignore") {
+        Some(numstr) => match from_str(numstr.as_slice()) {
+            Some(num) => num,
+            None => {
+                show_error!("\"{}\" is not a valid number", numstr);
+                return 1;
+            }
+        },
+        None => 0
+    };
+
+    if !matches.opt_present("all") {
+        ignore += match os::getenv("OMP_NUM_THREADS") {
+            Some(threadstr) => match from_str(threadstr.as_slice()) {
+                Some(num) => num,
+                None => 0
+            },
+            None => 0
+        };
+    }
+
+    let mut cores = os::num_cpus();
+    if cores <= ignore {
+        cores = 1;
+    } else {
+        cores -= ignore;
+    }
+    println!("{}", cores);
 
     return 0
 }