Resolve #14

riomucho/__init__.py

@@ -9,7 +9,7 @@
 
 import click
 import numpy as np
-import rasterio as rio
+import rasterio
 from rasterio.transform import guard_transform
 
 import riomucho.scripts.riomucho_utils as utils
@@ -72,7 +72,7 @@ def main_worker(inpaths, g_work_func, g_args):
     global srcs
     work_func = g_work_func
     global_args = g_args
-    srcs = [rio.open(i) for i in inpaths]
+    srcs = [rasterio.open(i) for i in inpaths]
 
 
 @job
@@ -102,11 +102,11 @@ class RioMucho:
     """TODO
     """
 
-    def __init__(self, inpaths, outpath, run_function, **kwargs):
+    def __init__(self, inpaths, outpath_or_dataset, run_function, **kwargs):
         """TODO
         """
         self.inpaths = inpaths
-        self.outpath = outpath
+        self.outpath_or_dataset = outpath_or_dataset
         self.run_function = run_function
 
         if 'mode' not in kwargs or kwargs['mode'] == 'simple_read':
@@ -155,9 +155,14 @@ def run(self, processes=4):
         else:
             reader_worker = simpleRead
 
+        if isinstance(self.outpath_or_dataset, rasterio.io.DatasetWriter):
+            destination = self.outpath_or_dataset
+        else:
+            destination = rasterio.open(self.outpath_or_dataset, 'w', **self.options)
+
         # Open an output file, work through the function in parallel,
         # and write out the data.
-        with rio.open(self.outpath, 'w', **self.options) as dst:
+        with destination as dst:
             for data, window in self.pool.imap_unordered(reader_worker, self.windows):
                 dst.write(data, window=window)
 

tests/test_mod.py

@@ -23,7 +23,7 @@ def test_riomucho_manual():
 
     with riomucho.RioMucho(['/tmp/test_1.tif', '/tmp/test_2.tif'], '/tmp/test_xyz_out.tif', read_function_manual,
         windows=windows,
-        global_args={}, 
+        global_args={},
         options=options,
         mode='manual_read') as rm:
 
@@ -75,8 +75,9 @@ def test_riomucho_arrayread():
 
 def test_riomucho_readmode_fail():
     with pytest.raises(ValueError):
-        with riomucho.RioMucho(['/tmp/test_1.tif',], '/tmp/test_xyz_out.tif', read_function_arrayread,
-            mode='mucho_gusto') as rm:
+        with riomucho.RioMucho(
+                ['/tmp/test_1.tif'], '/tmp/test_xyz_out.tif', read_function_arrayread,
+                mode='mucho_gusto') as rm:
             rm.run(4)
 
 
@@ -104,8 +105,10 @@ def test_pool_worker_exceptions(tmpdir):
     def fail(data, window, ij, g_args):
         return data * (1 / 0)
 
-    with riomucho.RioMucho(['/tmp/test_1.tif', '/tmp/test_2.tif'], str(tmpdir.join('output.tif')), fail,
-            mode='array_read', options=options) as rm:
+    with riomucho.RioMucho(
+            ['/tmp/test_1.tif', '/tmp/test_2.tif'],
+            str(tmpdir.join('output.tif')), fail, mode='array_read',
+            options=options) as rm:
         with pytest.raises(riomucho.MuchoChildError) as excinfo:
             rm.run(4)
 
@@ -120,5 +123,17 @@ def foo(*args, **kwargs):
 
     with pytest.raises(riomucho.MuchoChildError) as excinfo:
         foo()
-
     assert "ZeroDivisionError" in str(excinfo.value)
+
+
+def test_riomucho_simple_dataset_object(tmpdir):
+    """We can pass an open dataset for output"""
+    with rasterio.open('/tmp/test_1.tif') as src:
+        options = src.profile
+
+    with rasterio.open(str(tmpdir.join('output.tif')), 'w', **options) as dst:
+        with riomucho.RioMucho(['/tmp/test_1.tif'], dst, read_function_simple) as rm:
+            rm.run(1)
+
+    with rasterio.open(str(tmpdir.join('output.tif'))) as outputsrc:
+        assert numpy.sum(outputsrc.read(1)[:10, :10] != 0) == 0