cogeotiff / rio-tiler

User friendly Rasterio plugin to read raster datasets.
https://cogeotiff.github.io/rio-tiler/
BSD 3-Clause "New" or "Revised" License
502 stars 106 forks source link

enable dynamic definition of Reader for MultiBaseReader #711

Open vincentsarago opened 3 months ago

vincentsarago commented 3 months ago

The main PR goal is to add a method to enable dynamic definition of the Asset's Reader based on its type (e.g COG or NetCDF).

The addition is non-breaking, the default reader will always be self.reader

from typing import Type, Set

from rio_tiler.types import AssetInfo
import attr
from rio_tiler.io import STACReader as OfficialSTACReader
from rio_tiler.io import Reader, XarrayReader, BaseReader
valid_types = {
    "image/tiff; application=geotiff",
    "image/tiff; application=geotiff; profile=cloud-optimized",
    "image/tiff; profile=cloud-optimized; application=geotiff",
    "image/tiff; application=geotiff; profile=cloud-optimized",
    "image/vnd.stac.geotiff; cloud-optimized=true",
    "image/tiff",
    "image/x.geotiff",
    "image/jp2",
    "application/x-hdf5",
    "application/x-hdf",
    "application/vnd+zarr",
    "application/x-netcdf",
}

@attr.s
class STACReader(OfficialSTACReader):

    include_asset_types: Set[str] = attr.ib(default=valid_types)

    def _get_reader(self, asset_info: AssetInfo) -> Type[BaseReader]:
        """Get Asset Reader."""
        asset_type = asset_info.get("type", None)

        if asset_type and asset_type in [
            "application/x-hdf5",
            "application/x-hdf",
            "application/vnd.zarr",
            "application/x-netcdf",

        ]:
            return XarrayReader

        return Reader

with STACReader("https://planetarycomputer.microsoft.com/api/stac/v1/collections/noaa-cdr-sea-surface-temperature-optimum-interpolation/items/oisst-avhrr-v02r01.20240506") as src:
    print(src)
    print(src.assets)
    info = src._get_asset_info("netcdf")
    print(info)
    print(info["type"])
    print(src._get_reader(info))

>> STACReader(bounds=[-180, -90.0, 180, 90.0], crs=CRS.from_epsg(4326), assets=['err', 'ice', 'sst', 'anom', 'netcdf'], input='https://planetarycomputer.microsoft.com/api/stac/v1/collections/noaa-cdr-sea-surface-temperature-optimum-interpolation/items/oisst-avhrr-v02r01.20240506', item=<Item id=oisst-avhrr-v02r01.20240506>, tms=<TileMatrixSet title='Google Maps Compatible for the World' id='WebMercatorQuad' crs='[http://www.opengis.net/def/crs/EPSG/0/3857>](http://www.opengis.net/def/crs/EPSG/0/3857%3E), minzoom=0, maxzoom=24, geographic_crs=CRS.from_epsg(4326), include_assets=None, exclude_assets=None, exclude_asset_types=None, reader=<class 'rio_tiler.io.rasterio.Reader'>, reader_options={}, fetch_options={}, ctx=<class 'rasterio.env.Env'>, include_asset_types={'image/tiff; application=geotiff', 'image/jp2', 'application/x-hdf', 'image/tiff; application=geotiff; profile=cloud-optimized', 'image/x.geotiff', 'application/x-hdf5', 'image/vnd.stac.geotiff; cloud-optimized=true', 'application/x-netcdf', 'application/vnd+zarr', 'image/tiff; profile=cloud-optimized; application=geotiff', 'image/tiff'})
>> ['err', 'ice', 'sst', 'anom', 'netcdf']
>> {'url': 'https://noaacdr.blob.core.windows.net/sea-surface-temp-optimum-interpolation/data/v2.1/avhrr/202405/oisst-avhrr-v02r01.20240506.nc', 'metadata': {'created': '2024-05-21T09:12:00Z', 'updated': '2024-05-21T09:12:00Z'}, 'type': 'application/x-netcdf'}
>> application/x-netcdf
>> <class 'rio_tiler.io.xarray.XarrayReader'>

cc @abarciauskas-bgse