Repository: ungarj/mapchete
Pull request #473 (2022-06-28 14:57:20)

make raster read() functions conform with rasterios read() functions …
…in returning 2D or 3D arrays


Patch:
diff --git a/mapchete/formats/default/raster_file.py b/mapchete/formats/default/raster_file.py
index 161ebef..94e8278 100644
--- a/mapchete/formats/default/raster_file.py
+++ b/mapchete/formats/default/raster_file.py
@@ -265,6 +265,14 @@ class InputTile(base.InputTile):
         """
         Read reprojected & resampled input data.
 
+        Parameters
+        ----------
+        indexes : list or int
+            Either a list of band indexes or a single band index. If only a single
+            band index is given, the function returns a 2D array, otherwise a 3D array.
+        resampling : str
+            Resampling method to be used.
+
         Returns
         -------
         data : array
@@ -311,12 +319,9 @@ class InputTile(base.InputTile):
     def _get_band_indexes(self, indexes=None):
         """Return valid band indexes."""
         if indexes:
-            if isinstance(indexes, list):
-                return indexes
-            else:
-                return [indexes]
+            return indexes
         else:
-            return range(1, self.profile["count"] + 1)
+            return list(range(1, self.profile["count"] + 1))
 
 
 def get_segmentize_value(input_file=None, tile_pyramid=None):
diff --git a/mapchete/formats/default/tile_directory.py b/mapchete/formats/default/tile_directory.py
index d712f30..14edf44 100644
--- a/mapchete/formats/default/tile_directory.py
+++ b/mapchete/formats/default/tile_directory.py
@@ -290,7 +290,8 @@ class InputTile(base.InputTile):
             vector file: also run checks if reprojected geometry is valid,
             otherwise throw RuntimeError (default: True)
         indexes : list or int
-            raster file: a list of band numbers; None will read all.
+            Either a list of band indexes or a single band index. If only a single
+            band index is given, the function returns a 2D array, otherwise a 3D array.
         dst_nodata : int or float, optional
             raster file: if not set, the nodata value from the source dataset
             will be used
diff --git a/mapchete/io/raster.py b/mapchete/io/raster.py
index 6bf7863..c623405 100644
--- a/mapchete/io/raster.py
+++ b/mapchete/io/raster.py
@@ -68,7 +68,9 @@ def read_raster_window(
     tile : Tile
         a Tile object
     indexes : list or int
-        a list of band numbers; None will read all.
+        Either a list of band indexes or a single band index. If only a single
+        band index is given, the function returns a 2D array, otherwise a 3D array.
+        None will read all. This behavior mimicks rasterios.
     resampling : string
         one of "nearest", "average", "bilinear" or "lanczos"
     src_nodata : int or float, optional
@@ -131,7 +133,7 @@ def _read_raster_window(
             )
         dst_shape = (
             (len(indexes), tile.height, tile.width)
-            if len(indexes) > 1
+            if isinstance(indexes, list)
             else (tile.height, tile.width)
         )
         return ma.masked_array(
@@ -177,13 +179,10 @@ def _read_raster_window(
         try:
             input_file = input_files
             dst_shape = tile.shape
-
             if not isinstance(indexes, int):
                 if indexes is None:
                     dst_shape = (None,) + dst_shape
-                elif len(indexes) == 1:
-                    indexes = indexes[0]
-                else:
+                elif isinstance(indexes, list):
                     dst_shape = (len(indexes),) + dst_shape
             # Check if potentially tile boundaries exceed tile matrix boundaries on
             # the antimeridian, the northern or the southern boundary.
diff --git a/mapchete/processes/contours.py b/mapchete/processes/contours.py
index b6010b2..5571ce8 100644
--- a/mapchete/processes/contours.py
+++ b/mapchete/processes/contours.py
@@ -77,6 +77,7 @@ def execute(
     ) as dem:
         logger.debug("reading input raster")
         dem_data = dem.read(
+            1,
             resampling=resampling,
             matching_method=td_matching_method,
             matching_max_zoom=td_matching_max_zoom,


Test patch:
diff --git a/test/test_commons.py b/test/test_commons.py
index 1a2949a..cca7def 100644
--- a/test/test_commons.py
+++ b/test/test_commons.py
@@ -67,7 +67,7 @@ def test_contours(cleantopo_tl):
             input=mp.config.get_inputs_for_tile(tile),
         )
         with user_process.open("file1") as dem:
-            arr = dem.read()
+            arr = dem.read(1)
             # valid contours
             contours = user_process.contours(arr)
             assert contours
diff --git a/test/test_formats.py b/test/test_formats.py
index ababd61..c1698c8 100644
--- a/test/test_formats.py
+++ b/test/test_formats.py
@@ -148,7 +148,7 @@ def test_read_from_raster_file(cleantopo_br):
             input=mp.config.get_inputs_for_tile(tile),
         )
         with user_process.open("file1") as f:
-            assert f.read().shape == f.read([1]).shape == f.read(1).shape
+            assert f.read().shape == f.read([1]).shape == (1, *f.read(1).shape)
 
 
 def test_invalid_input_type(example_mapchete):
diff --git a/test/test_formats_raster_file.py b/test/test_formats_raster_file.py
new file mode 100644
index 0000000..1cb0624
--- /dev/null
+++ b/test/test_formats_raster_file.py
@@ -0,0 +1,18 @@
+from mapchete.formats.default.raster_file import InputData
+
+
+def test_read_indexes_shape(cleantopo_br_tiledir, cleantopo_br_tif):
+    input_data = InputData({"path": cleantopo_br_tif})
+    input_tile = input_data.open(cleantopo_br_tiledir.first_process_tile())
+
+    # no indexes --> 3D array
+    three_d_arr = input_tile.read()
+    assert three_d_arr.ndim == 3
+
+    # list index --> 3D array
+    three_d_arr = input_tile.read([1])
+    assert three_d_arr.ndim == 3
+
+    # int index --> 2D array
+    two_d_arr = input_tile.read(1)
+    assert two_d_arr.ndim == 2
diff --git a/test/test_formats_tiledir_input.py b/test/test_formats_tiledir_input.py
index c43bb2a..2bc2960 100644
--- a/test/test_formats_tiledir_input.py
+++ b/test/test_formats_tiledir_input.py
@@ -6,6 +6,7 @@ import shutil
 
 import mapchete
 from mapchete.formats import available_input_formats
+from mapchete.formats.default.tile_directory import InputData
 from mapchete.errors import MapcheteDriverError
 
 
@@ -149,6 +150,26 @@ def test_read_from_dir(mp_tmpdir, cleantopo_br, cleantopo_br_tiledir):
     _run_tiledir_process_raster(config, 4, bounds)
 
 
+def test_read_indexes_shape(cleantopo_br_tiledir):
+    mp = cleantopo_br_tiledir.mp()
+    # create local TileDirectory
+    list(mp.compute())
+    input_data = InputData({"path": mp.config.output_reader.path})
+    input_tile = input_data.open(cleantopo_br_tiledir.first_process_tile())
+
+    # no indexes --> 3D array
+    three_d_arr = input_tile.read()
+    assert three_d_arr.ndim == 3
+
+    # list index --> 3D array
+    three_d_arr = input_tile.read([1])
+    assert three_d_arr.ndim == 3
+
+    # int index --> 2D array
+    two_d_arr = input_tile.read(1)
+    assert two_d_arr.ndim == 2
+
+
 def test_no_metadata_json(mp_tmpdir, cleantopo_br_tiledir):
     """Read raster data."""
     # prepare data
diff --git a/test/test_io_raster.py b/test/test_io_raster.py
index ab46b00..75ccda8 100644
--- a/test/test_io_raster.py
+++ b/test/test_io_raster.py
@@ -86,7 +86,7 @@ def test_read_raster_window_reproject(dummy1_3857_tif, minmax_zoom):
             assert band.shape == tile.shape
         bands = read_raster_window(dummy1_3857_tif, tile, [1])
         assert isinstance(bands, ma.MaskedArray)
-        assert bands.shape == tile.shape
+        assert bands.shape == (1, *tile.shape)
     # errors
     with pytest.raises(IOError):
         read_raster_window("nonexisting_path", tile)