Rio_slurp can now reproject as well

If given GeoBox call rerpoject instead of read

datacube/testutils/io.py

@@ -127,7 +127,40 @@ def rio_geobox(meta):
     return GeoBox(w, h, transform, crs)
 
 
-def rio_slurp(fname, out_shape=None, **kw):
+def rio_slurp_reproject(fname, gbox, dtype=None, dst_nodata=None, **kw):
+    """
+    Read image with reprojection
+    """
+    import rasterio
+    from rasterio.warp import reproject
+
+    with rasterio.open(str(fname), 'r') as src:
+        src_band = rasterio.band(src, 1)
+
+        if dtype is None:
+            dtype = src.dtypes[0]
+        if dst_nodata is None:
+            dst_nodata = src.nodata
+        if dst_nodata is None:
+            dst_nodata = 0
+
+        pix = np.full(gbox.shape, dst_nodata, dtype=dtype)
+
+        reproject(src_band, pix,
+                  dst_nodata=dst_nodata,
+                  dst_transform=gbox.transform,
+                  dst_crs=str(gbox.crs),
+                  **kw)
+
+        meta = src.meta
+        meta['src_gbox'] = rio_geobox(meta)
+        meta['path'] = fname
+        meta['gbox'] = gbox
+
+        return pix, meta
+
+
+def rio_slurp_read(fname, out_shape=None, **kw):
     """
     Read whole image file using rasterio.
 
@@ -144,3 +177,23 @@ def rio_slurp(fname, out_shape=None, **kw):
         meta['gbox'] = rio_geobox(meta)
         meta['path'] = fname
         return data, meta
+
+
+def rio_slurp(fname, *args, **kw):
+    """
+    Dispatches to either:
+
+    rio_slurp_read(fname, out_shape, ..)
+    rio_slurp_reproject(fname, gbox, ...)
+
+    """
+    if len(args) == 0:
+        if 'gbox' in kw:
+            return rio_slurp_reproject(fname, **kw)
+        else:
+            return rio_slurp_read(fname, **kw)
+
+    if isinstance(args[0], GeoBox):
+        return rio_slurp_reproject(fname, *args, **kw)
+    else:
+        return rio_slurp_read(fname, *args, **kw)

tests/storage/test_storage_read.py

@@ -17,6 +17,10 @@
 
 from datacube.utils.geometry import gbox as gbx
 
+from datacube.testutils.io import (
+    rio_slurp
+)
+
 from datacube.testutils.geom import (
     epsg3857,
     AlbersGS,
@@ -187,6 +191,9 @@ def _read(gbox,
     assert roi[0].start > 0 and roi[1].start > 0
     assert (yy[0] == -999).all()
 
+    yy_expect, _ = rio_slurp(mm.path, gbox)
+    np.testing.assert_array_equal(yy, yy_expect)
+
     gbox = gbx.zoom_out(mm.gbox[3:-3, 10:-10], 2.1)
     yy, roi = _read(gbox)
 

tests/test_load_data.py

@@ -9,7 +9,7 @@
     mk_sample_dataset,
     mk_test_image,
 )
-from datacube.testutils.io import write_gtiff
+from datacube.testutils.io import write_gtiff, rio_slurp
 
 
 def gen_tiff_dataset(bands,
@@ -97,3 +97,36 @@ def custom_fuser(dest, delta):
     assert ds_data.aa.nodata == nodata
     assert custom_fuser_call_count > 0
     np.testing.assert_array_equal(nodata + aa + aa, ds_data.aa.values[0])
+
+
+def test_rio_slurp(tmpdir):
+    w, h, dtype, nodata, ndw = 96, 64, 'int16', -999, 7
+
+    pp = Path(str(tmpdir))
+
+    aa = mk_test_image(w, h, dtype, nodata, nodata_width=ndw)
+
+    assert aa.shape == (h, w)
+    assert aa.dtype.name == dtype
+    assert aa[10, 30] == (30 << 8) | 10
+    assert aa[10, 11] == nodata
+
+    aa0 = aa.copy()
+    mm = write_gtiff(pp/"rio-slurp-aa.tif", aa, nodata=-999, overwrite=True)
+    mm = SimpleNamespace(**mm)
+
+    aa, _ = rio_slurp(mm.path)
+    np.testing.assert_array_equal(aa, aa0)
+
+    aa, _ = rio_slurp(mm.path, aa0.shape)
+    np.testing.assert_array_equal(aa, aa0)
+
+    aa, _ = rio_slurp(mm.path, mm.gbox)
+    np.testing.assert_array_equal(aa, aa0)
+
+    aa, _ = rio_slurp(mm.path, gbox=mm.gbox, dtype='float32')
+    assert aa.dtype == 'float32'
+    np.testing.assert_array_equal(aa, aa0.astype('float32'))
+
+    aa, _ = rio_slurp(mm.path, mm.gbox, dst_nodata=-33)
+    np.testing.assert_array_equal(aa == -33, aa0 == -999)