Repository: mapbox/rio-color

Issue:
use all cores with -j -1
Modify max_procs so that a negative value means "use the number of cores available on this machine"

Comments:

Patch:
diff --git a/rio_color/scripts/cli.py b/rio_color/scripts/cli.py
index 7ec9cab..9580f68 100755
--- a/rio_color/scripts/cli.py
+++ b/rio_color/scripts/cli.py
@@ -1,13 +1,27 @@
 import click
+
 import rasterio
 from rio_color.workers import atmos_worker, color_worker
 from rio_color.operations import parse_operations
 import riomucho
 
 
+jobs_opt = click.option(
+    '--jobs', '-j', type=int, default=1,
+    help="Number of jobs to run simultaneously, Use -1 for all cores, default: 1")
+
+
+def check_jobs(jobs):
+    if jobs == 0:
+        raise click.UsageError("Jobs must be >= 1 or == -1")
+    elif jobs < 0:
+        import multiprocessing
+        jobs = multiprocessing.cpu_count()
+    return jobs
+
+
 @click.command('color')
-@click.option('--jobs', '-j', type=int, default=1,
-              help="Number of jobs to run simultaneously, default: 1")
+@jobs_opt
 @click.option('--out-dtype', '-d', type=click.Choice(['uint8', 'uint16']),
               help="Integer data type for output data, default: same as input")
 @click.argument('src_path', type=click.Path(exists=True))
@@ -73,6 +87,8 @@ Example:
         'out_dtype': out_dtype
     }
 
+    jobs = check_jobs(jobs)
+
     if jobs > 1:
         with riomucho.RioMucho(
             [src_path],
@@ -103,8 +119,7 @@ Example:
 @click.option('--bias', '-b', type=click.FLOAT, default=15,
               help="Skew (brighten/darken) the output. Lower values make it "
                    "brighter. 0..100 (50 is none), default: 15.")
-@click.option('--jobs', '-j', type=int, default=1,
-              help="Number of jobs to run simultaneously, default: 1")
+@jobs_opt
 @click.option('--out-dtype', '-d', type=click.Choice(['uint8', 'uint16']),
               help="Integer data type for output data, default: same as input")
 @click.argument('src_path', type=click.Path(exists=True))
@@ -132,6 +147,8 @@ def atmos(ctx, atmo, contrast, bias, jobs, out_dtype,
         'out_dtype': out_dtype
     }
 
+    jobs = check_jobs(jobs)
+
     if jobs > 1:
         with riomucho.RioMucho(
             [src_path],


Test patch:
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 8448a74..1d30644 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -1,10 +1,12 @@
 import os
 
+from click import UsageError
 from click.testing import CliRunner
 import numpy as np
+import pytest
 import rasterio
 
-from rio_color.scripts.cli import color, atmos
+from rio_color.scripts.cli import color, atmos, check_jobs
 
 
 def equal(r1, r2):
@@ -98,3 +100,25 @@ def test_bad_op(tmpdir):
     assert result.exit_code == 2
     assert "foob is not a valid operation" in result.output
     assert not os.path.exists(output)
+
+
+def test_color_jobsn1(tmpdir):
+    output = str(tmpdir.join('colorj1.tif'))
+    runner = CliRunner()
+    result = runner.invoke(
+        color,
+        [
+            '-d', 'uint8',
+            '-j', '-1',
+            'tests/rgb8.tif',
+            output,
+            "gamma 1,2,3 1.85"])
+    assert result.exit_code == 0
+    assert os.path.exists(output)
+
+
+def test_check_jobs():
+    assert 1 == check_jobs(1)
+    assert check_jobs(-1) > 0
+    with pytest.raises(UsageError):
+        check_jobs(0)