(backport) avoid indexes collision for multi-base Reader

cogeotiff · Nov 15, 2023 · c069da4 · c069da4
1 parent a814aaf
commit c069da4
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 6 deletions.
diff --git a/CHANGES.md b/CHANGES.md
@@ -1,5 +1,4 @@
 
-=======
 # 6.2.3.post1 (2023-11-16)
 
 * validate `shape` in `ImageData.get_coverage_array` to avoid rasterio error when re-projecting the geometry [Backported from 6.2.6]

diff --git a/rio_tiler/io/base.py b/rio_tiler/io/base.py
@@ -489,8 +489,11 @@ def tile(
 
         asset_indexes = asset_indexes or {}
 
+        # We fall back to `indexes` if provided
+        indexes = kwargs.pop("indexes", None)
+
         def _reader(asset: str, *args: Any, **kwargs: Any) -> ImageData:
-            idx = asset_indexes.get(asset) or kwargs.pop("indexes", None)  # type: ignore
+            idx = asset_indexes.get(asset) or indexes  # type: ignore
 
             asset_info = self._get_asset_info(asset)
             url = asset_info["url"]
@@ -567,8 +570,11 @@ def part(
 
         asset_indexes = asset_indexes or {}
 
+        # We fall back to `indexes` if provided
+        indexes = kwargs.pop("indexes", None)
+
         def _reader(asset: str, *args: Any, **kwargs: Any) -> ImageData:
-            idx = asset_indexes.get(asset) or kwargs.pop("indexes", None)  # type: ignore
+            idx = asset_indexes.get(asset) or indexes  # type: ignore
 
             asset_info = self._get_asset_info(asset)
             url = asset_info["url"]
@@ -643,8 +649,11 @@ def preview(
 
         asset_indexes = asset_indexes or {}
 
+        # We fall back to `indexes` if provided
+        indexes = kwargs.pop("indexes", None)
+
         def _reader(asset: str, **kwargs: Any) -> ImageData:
-            idx = asset_indexes.get(asset) or kwargs.pop("indexes", None)  # type: ignore
+            idx = asset_indexes.get(asset) or indexes  # type: ignore
 
             asset_info = self._get_asset_info(asset)
             url = asset_info["url"]
@@ -723,8 +732,11 @@ def point(
 
         asset_indexes = asset_indexes or {}
 
+        # We fall back to `indexes` if provided
+        indexes = kwargs.pop("indexes", None)
+
         def _reader(asset: str, *args, **kwargs: Any) -> PointData:
-            idx = asset_indexes.get(asset) or kwargs.pop("indexes", None)  # type: ignore
+            idx = asset_indexes.get(asset) or indexes  # type: ignore
 
             asset_info = self._get_asset_info(asset)
             url = asset_info["url"]
@@ -795,8 +807,11 @@ def feature(
 
         asset_indexes = asset_indexes or {}
 
+        # We fall back to `indexes` if provided
+        indexes = kwargs.pop("indexes", None)
+
         def _reader(asset: str, *args: Any, **kwargs: Any) -> ImageData:
-            idx = asset_indexes.get(asset) or kwargs.pop("indexes", None)  # type: ignore
+            idx = asset_indexes.get(asset) or indexes  # type: ignore
 
             asset_info = self._get_asset_info(asset)
             url = asset_info["url"]

diff --git a/tests/test_io_stac.py b/tests/test_io_stac.py
@@ -199,6 +199,22 @@ def test_tile_valid(rio):
         assert img.mask.shape == (256, 256)
         assert img.band_names == ["green_b1", "green_b1", "red_b1", "red_b1"]
 
+        # check that indexes and asset_indexes are not conflicting
+        img = stac.tile(
+            71,
+            102,
+            8,
+            assets=("green", "red"),
+            indexes=None,
+            asset_indexes={
+                "green": (1,),
+                "red": 1,
+            },
+        )
+        assert img.data.shape == (2, 256, 256)
+        assert img.mask.shape == (256, 256)
+        assert img.band_names == ["green_b1", "red_b1"]
+
         img = stac.tile(71, 102, 8, expression="green_b1*2;green_b1;red_b1*2")
         assert img.data.shape == (3, 256, 256)
         assert img.mask.shape == (256, 256)