SubDatasetsReader.py

@attr.s
class SDSTReader(COGReader):
    """Cloud Optimized GeoTIFF Reader."""

    subdatasets: Dict[int, str] = attr.ib(init=False)

    def __attrs_post_init__(self):
        """Define _kwargs, open dataset and get info."""
        super().__attrs_post_init__()

        self.subdatasets = {ix+1: val for ix, val in enumerate(self.dataset.subdatasets)}

    def info(self) -> Info:
        cog_info = super().info()
        cog_info.subdatasets = self.subdatasets
        return cog_info

    def statistics(
        self,
        categorical: bool = False,
        categories: Optional[List[float]] = None,
        percentiles: List[int] = [2, 98],
        hist_options: Optional[Dict] = None,
        max_size: int = 1024,
        **kwargs: Any,
    ) -> Dict[str, BandStatistics]:
        raise NotImplementedError("nope")

    def tile(
        self,
        tile_x: int,
        tile_y: int,
        tile_z: int,
        sdst: Optional[Indexes] = None,
        **kwargs: Any,
    ) -> ImageData:

        sdst = sdst or list(self.subdatasets)

        if isinstance(sdst, str):
            sdst = (sdst,)

        def _reader(idx: int, *args: Any, **kwargs: Any) -> ImageData:
            with COGReader(self.subdatasets[idx], tms=self.tms, **self._kwargs) as cog:  # type: ignore
                data = cog.tile(*args, **kwargs)
                data.band_names = [f"{idx}_{n}" for n in data.band_names]
                return data

        return ImageData.create_from_list(
            [_reader(idx,  tile_x, tile_y, tile_z,  **kwargs) for idx in sdst]
        )

    def part(
        self,
        bbox: BBox,
        dst_crs: Optional[CRS] = None,
        bounds_crs: CRS = WGS84_CRS,
        indexes: Optional[Union[int, Sequence]] = None,
        expression: Optional[str] = None,
        max_size: Optional[int] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        **kwargs: Any,
    ) -> ImageData:
        pass

    def preview(
        self, sdst: Optional[Indexes] = None, **kwargs: Any,
    ) -> ImageData:

        sdst = sdst or list(self.subdatasets)

        if isinstance(sdst, str):
            sdst = (sdst,)

        def _reader(idx: int, **kwargs: Any) -> ImageData:
            with COGReader(self.subdatasets[idx], **self._kwargs) as cog:
                data = cog.preview(**kwargs)
                data.band_names = [f"{idx}_{n}" for n in data.band_names]
                return data

        return ImageData.create_from_list([_reader(idx, **kwargs) for idx in sdst])

    def point(
        self, lon: float, lat: float, sdst: Optional[Indexes] = None, **kwargs: Any,
    ) -> List:
        sdst = sdst or list(self.subdatasets)

        if isinstance(sdst, str):
            sdst = (sdst,)

        def _reader(idx: int, *args: Any, **kwargs: Any) -> ImageData:
            with COGReader(self.subdatasets[idx], **self._kwargs) as cog:
                return cog.point(*args, **kwargs)

        return [_reader(idx, lon, lat, **kwargs) for idx in sdst]

    def feature(
        self, shape: Dict, sdst: Optional[Indexes] = None, **kwargs: Any,
    ) -> ImageData:

        sdst = sdst or list(self.subdatasets)

        if isinstance(sdst, str):
            sdst = (sdst,)

        def _reader(idx: int, *args: Any, **kwargs: Any) -> ImageData:
            with COGReader(self.subdatasets[idx], **self._kwargs) as cog:
                data = cog.feature(*args, **kwargs)
                data.band_names = [f"{idx}_{n}" for n in data.band_names]
                return data

        return ImageData.create_from_list([_reader(idx, shape, **kwargs) for idx in sdst])

    def read(
        self, sdst: Optional[Indexes] = None, **kwargs: Any,
    ) -> ImageData:
        sdst = sdst or list(self.subdatasets)

        if isinstance(sdst, str):
            sdst = (sdst,)

        def _reader(idx: int, **kwargs: Any) -> ImageData:
            with COGReader(self.subdatasets[idx], **self._kwargs) as cog:
                data = cog.read(**kwargs)
                data.band_names = [f"{idx}_{n}" for n in data.band_names]
                return data

        return ImageData.create_from_list([_reader(idx, **kwargs) for idx in sdst])