Merge pull request #19 from developmentseed/features/allow-file-separ…

…ated-bands

allow file separated band storage

titiler/cmr/backend.py

@@ -1,6 +1,8 @@
 """TiTiler.cmr custom Mosaic Backend."""
 
-from typing import Any, Dict, List, Optional, Tuple, Type, TypedDict
+import os
+import re
+from typing import Any, Dict, List, Optional, Tuple, Type, TypedDict, Union
 
 import attr
 import earthaccess
@@ -40,8 +42,7 @@ def aws_s3_credential(auth: Auth, provider: str) -> Dict:
 class Asset(TypedDict, total=False):
     """Simple Asset model."""
 
-    url: str
-    type: str
+    url: Union[str, Dict[str, str]]
     provider: str
 
 
@@ -160,6 +161,7 @@ def get_assets(
         xmax: float,
         ymax: float,
         limit: int = 100,
+        bands_regex: Optional[str] = None,
         **kwargs: Any,
     ) -> List[Asset]:
         """Find assets."""
@@ -172,14 +174,30 @@ def get_assets(
 
         assets: List[Asset] = []
         for r in results:
-            assets.append(
-                {
-                    "url": r.data_links(access="direct")[
-                        0
-                    ],  # NOTE: should we not do this?
-                    "provider": r["meta"]["provider-id"],
-                }
-            )
+            if bands_regex:
+                links = r.data_links(access="direct")
+
+                band_urls = []
+                for url in links:
+                    if match := re.search(bands_regex, os.path.basename(url)):
+                        band_urls.append((match.group(), url))
+
+                urls = dict(band_urls)
+                if urls:
+                    assets.append(
+                        {
+                            "url": urls,
+                            "provider": r["meta"]["provider-id"],
+                        }
+                    )
+
+            else:
+                assets.append(
+                    {
+                        "url": r.data_links(access="direct")[0],
+                        "provider": r["meta"]["provider-id"],
+                    }
+                )
 
         return assets
 
@@ -193,6 +211,7 @@ def tile(
         tile_y: int,
         tile_z: int,
         cmr_query: Dict,
+        bands_regex: Optional[str] = None,
         **kwargs: Any,
     ) -> Tuple[ImageData, List[str]]:
         """Get Tile from multiple observation."""
@@ -201,6 +220,7 @@ def tile(
             tile_y,
             tile_z,
             **cmr_query,
+            bands_regex=bands_regex,
         )
 
         if not mosaic_assets:
@@ -259,6 +279,7 @@ def point(
         lat: float,
         cmr_query: Dict,
         coord_crs: CRS = WGS84_CRS,
+        bands_regex: Optional[str] = None,
         **kwargs: Any,
     ) -> List:
         """Get Point value from multiple observation."""
@@ -270,6 +291,7 @@ def part(
         cmr_query: Dict,
         dst_crs: Optional[CRS] = None,
         bounds_crs: CRS = WGS84_CRS,
+        bands_regex: Optional[str] = None,
         **kwargs: Any,
     ) -> Tuple[ImageData, List[str]]:
         """Create an Image from multiple items for a bbox."""
@@ -282,6 +304,7 @@ def feature(
         dst_crs: Optional[CRS] = None,
         shape_crs: CRS = WGS84_CRS,
         max_size: int = 1024,
+        bands_regex: Optional[str] = None,
         **kwargs: Any,
     ) -> Tuple[ImageData, List[str]]:
         """Create an Image from multiple items for a GeoJSON feature."""

titiler/cmr/factory.py

@@ -3,7 +3,7 @@
 import json
 import re
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Literal, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, Type, Union
 from urllib.parse import urlencode
 
 import jinja2
@@ -14,7 +14,7 @@
 from morecantile import tms as default_tms
 from morecantile.defaults import TileMatrixSets
 from pydantic import conint
-from rio_tiler.io import Reader
+from rio_tiler.io import BaseReader, Reader
 from rio_tiler.types import RIOResampling, WarpResampling
 from starlette.requests import Request
 from starlette.responses import HTMLResponse, Response
@@ -26,7 +26,7 @@
 from titiler.cmr.backend import CMRBackend
 from titiler.cmr.dependencies import OutputType, cmr_query
 from titiler.cmr.enums import MediaType
-from titiler.cmr.reader import ZarrReader
+from titiler.cmr.reader import MultiFilesBandsReader, ZarrReader
 from titiler.core import dependencies
 from titiler.core.algorithm import algorithms as available_algorithms
 from titiler.core.factory import img_endpoint_params
@@ -451,7 +451,7 @@ def tiles_endpoint(
                 ),
             ] = None,
             ###################################################################
-            # COG Reader Options
+            # Rasterio Reader Options
             ###################################################################
             indexes: Annotated[
                 Optional[List[int]],
@@ -468,6 +468,20 @@ def tiles_endpoint(
                     description="rio-tiler's band math expression",
                 ),
             ] = None,
+            bands: Annotated[
+                Optional[List[str]],
+                Query(
+                    title="Band names",
+                    description="Band names.",
+                ),
+            ] = None,
+            bands_regex: Annotated[
+                Optional[str],
+                Query(
+                    title="Regex expression to parse dataset links",
+                    description="Regex expression to parse dataset links.",
+                ),
+            ] = None,
             unscale: Annotated[
                 Optional[bool],
                 Query(
@@ -516,10 +530,12 @@ def tiles_endpoint(
 
             tms = self.supported_tms.get(tileMatrixSetId)
 
-            read_options: Dict[str, Any] = {}
-            reader_options: Dict[str, Any] = {}
+            read_options: Dict[str, Any]
+            reader_options: Dict[str, Any]
+            options: Dict[str, Any]
+            reader: Type[BaseReader]
 
-            if backend != "cog":
+            if backend != "rasterio":
                 reader = ZarrReader
                 read_options = {}
 
@@ -531,16 +547,34 @@ def tiles_endpoint(
                 }
                 reader_options = {k: v for k, v in options.items() if v is not None}
             else:
-                reader = Reader
-                options = {
-                    "indexes": indexes,  # type: ignore
-                    "expression": expression,
-                    "unscale": unscale,
-                    "resampling_method": resampling_method,
-                }
-                read_options = {k: v for k, v in options.items() if v is not None}
-
-                reader_options = {}
+                if bands_regex:
+                    assert (
+                        bands
+                    ), "`bands=` option must be provided when using Multi bands collections."
+
+                    reader = MultiFilesBandsReader
+                    options = {
+                        "expression": expression,
+                        "bands": bands,
+                        "unscale": unscale,
+                        "resampling_method": resampling_method,
+                        "bands_regex": bands_regex,
+                    }
+                    read_options = {k: v for k, v in options.items() if v is not None}
+                    reader_options = {}
+
+                else:
+                    assert bands, "Can't use `bands=` option without `bands_regex`"
+
+                    reader = Reader
+                    options = {
+                        "indexes": indexes,
+                        "expression": expression,
+                        "unscale": unscale,
+                        "resampling_method": resampling_method,
+                    }
+                    read_options = {k: v for k, v in options.items() if v is not None}
+                    reader_options = {}
 
             with CMRBackend(
                 tms=tms,

titiler/cmr/reader.py

@@ -6,7 +6,7 @@
 import contextlib
 import pickle
 import re
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Type
 
 import attr
 import fsspec
@@ -17,6 +17,8 @@
 from morecantile import TileMatrixSet
 from rasterio.crs import CRS
 from rio_tiler.constants import WEB_MERCATOR_TMS, WGS84_CRS
+from rio_tiler.errors import InvalidBandName
+from rio_tiler.io import BaseReader, MultiBandReader, Reader
 from rio_tiler.io.xarray import XarrayReader
 from rio_tiler.types import BBox
 
@@ -284,3 +286,45 @@ def list_variables(
             consolidated=consolidated,
         ) as ds:
             return list(ds.data_vars)  # type: ignore
+
+
+@attr.s
+class MultiFilesBandsReader(MultiBandReader):
+    """Multiple Files as Bands."""
+
+    input: Dict[str, str] = attr.ib()
+    tms: TileMatrixSet = attr.ib(default=WEB_MERCATOR_TMS)
+
+    reader_options: Dict = attr.ib(factory=dict)
+    reader: Type[BaseReader] = attr.ib(default=Reader)
+
+    minzoom: int = attr.ib()
+    maxzoom: int = attr.ib()
+
+    @minzoom.default
+    def _minzoom(self):
+        return self.tms.minzoom
+
+    @maxzoom.default
+    def _maxzoom(self):
+        return self.tms.maxzoom
+
+    def __attrs_post_init__(self):
+        """Fetch Reference band to get the bounds."""
+        self.bands = list(self.input)
+        # with self.reader(
+        #     self.input[0],
+        #     tms=self.tms,
+        #     **self.reader_options,
+        # ) as cog:
+        #     self.bounds = cog.bounds
+        #     self.crs = cog.crs
+        #     self.minzoom = cog.minzoom
+        #     self.maxzoom = cog.maxzoom
+
+    def _get_band_url(self, band: str) -> str:
+        """Validate band's name and return band's url."""
+        if band not in self.bands:
+            raise InvalidBandName(f"{band} is not valid")
+
+        return self.input[band]