rio-tiler 7.9.2__py3-none-any.whl → 8.0.0__py3-none-any.whl

rio_tiler/io/xarray.py

@@ -3,8 +3,9 @@
 from __future__ import annotations
 
 import math
+import os
 import warnings
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Tuple
 
 import attr
 import numpy
@@ -23,6 +24,7 @@ from rasterio.warp import transform_bounds, transform_geom
 from rio_tiler.constants import WEB_MERCATOR_CRS, WEB_MERCATOR_TMS, WGS84_CRS
 from rio_tiler.errors import (
     InvalidGeographicBounds,
+    MaxArraySizeError,
     MissingCRS,
     PointOutsideBounds,
     TileOutsideBounds,
@@ -50,6 +52,9 @@ except ImportError:  # pragma: nocover
     rioxarray = None  # type: ignore
 
 
+MAX_ARRAY_SIZE = os.environ.get("RIO_TILER_MAX_ARRAY_SIZE", 1_000_000_000)  # 1Gb
+
+
 @attr.s
 class XarrayReader(BaseReader):
     """Xarray Reader.
@@ -133,7 +138,7 @@ class XarrayReader(BaseReader):
         return self._maxzoom
 
     @property
-    def band_names(self) -> List[str]:
+    def band_descriptions(self) -> List[str]:
         """
         Return list of `band descriptions` in DataArray.
 
@@ -155,7 +160,7 @@ class XarrayReader(BaseReader):
             if coords_name:
                 return [str(self.input.coords[coords_name[0]].data)]
 
-            return [self.input.name or "array"]
+            return [self.input.name or ""]
 
         return [str(band) for d in self._dims for band in self.input[d].values]
 
@@ -168,7 +173,7 @@ class XarrayReader(BaseReader):
             "crs": CRS_to_uri(self.crs) or self.crs.to_wkt(),
             "band_metadata": [(f"b{ix}", v) for ix, v in enumerate(metadata, 1)],
             "band_descriptions": [
-                (f"b{ix}", v) for ix, v in enumerate(self.band_names, 1)
+                (f"b{ix}", v) for ix, v in enumerate(self.band_descriptions, 1)
             ],
             "dtype": str(self.input.dtype),
             "nodata_type": "Nodata" if self.input.rio.nodata is not None else "None",
@@ -186,11 +191,14 @@ class XarrayReader(BaseReader):
         return Info(**meta)
 
     def _sel_indexes(
-        self, indexes: Optional[Indexes] = None
-    ) -> Tuple[xarray.DataArray, List[str]]:
+        self,
+        indexes: Indexes | None = None,
+    ) -> Tuple[xarray.DataArray, List[str], List[str]]:
         """Select `band` indexes in DataArray."""
         da = self.input
-        band_names = self.band_names
+        band_descriptions = self.band_descriptions
+        band_names = [f"b{ix + 1}" for ix in range(self.input.rio.count)]
+
         if indexes := cast_to_sequence(indexes):
             assert all(v > 0 for v in indexes), "Indexes value must be >= 1"
             if da.ndim == 2:
@@ -199,28 +207,30 @@ class XarrayReader(BaseReader):
                         f"Invalid indexes {indexes} for array of shape {da.shape}"
                     )
 
-                return da, band_names
+                return da, band_names, band_descriptions
 
             indexes = [idx - 1 for idx in indexes]
+
             da = da[indexes]
-            band_names = [self.band_names[idx] for idx in indexes]
+            band_descriptions = [band_descriptions[idx] for idx in indexes]
+            band_names = [band_names[idx] for idx in indexes]
 
-        return da, band_names
+        return da, band_names, band_descriptions
 
     def statistics(
         self,
         categorical: bool = False,
-        categories: Optional[List[float]] = None,
-        percentiles: Optional[List[int]] = None,
-        hist_options: Optional[Dict] = None,
-        nodata: Optional[NoData] = None,
-        indexes: Optional[Indexes] = None,
+        categories: List[float] | None = None,
+        percentiles: List[int] | None = None,
+        hist_options: Dict | None = None,
+        nodata: NoData | None = None,
+        indexes: Indexes | None = None,
         **kwargs: Any,
     ) -> Dict[str, BandStatistics]:
         """Return statistics from a dataset."""
         hist_options = hist_options or {}
 
-        da, band_names = self._sel_indexes(indexes)
+        da, band_names, _ = self._sel_indexes(indexes)
 
         if nodata is not None:
             da = da.rio.write_nodata(nodata)
@@ -246,8 +256,8 @@ class XarrayReader(BaseReader):
         tilesize: int = 256,
         reproject_method: WarpResampling = "nearest",
         auto_expand: bool = True,
-        nodata: Optional[NoData] = None,
-        indexes: Optional[Indexes] = None,
+        nodata: NoData | None = None,
+        indexes: Indexes | None = None,
         out_dtype: str | numpy.dtype | None = None,
         **kwargs: Any,
     ) -> ImageData:
@@ -291,15 +301,15 @@ class XarrayReader(BaseReader):
     def part(  # noqa: C901
         self,
         bbox: BBox,
-        dst_crs: Optional[CRS] = None,
+        dst_crs: CRS | None = None,
         bounds_crs: CRS = WGS84_CRS,
         reproject_method: WarpResampling = "nearest",
         auto_expand: bool = True,
-        nodata: Optional[NoData] = None,
-        indexes: Optional[Indexes] = None,
-        max_size: Optional[int] = None,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
+        nodata: NoData | None = None,
+        indexes: Indexes | None = None,
+        max_size: int | None = None,
+        height: int | None = None,
+        width: int | None = None,
         resampling_method: RIOResampling = "nearest",
         out_dtype: str | numpy.dtype | None = None,
         **kwargs: Any,
@@ -331,7 +341,7 @@ class XarrayReader(BaseReader):
 
         dst_crs = dst_crs or bounds_crs
 
-        da, band_names = self._sel_indexes(indexes)
+        da, band_names, band_descriptions = self._sel_indexes(indexes)
 
         if nodata is not None:
             da = da.rio.write_nodata(nodata)
@@ -348,6 +358,11 @@ class XarrayReader(BaseReader):
             auto_expand=auto_expand,
         )
 
+        if da.nbytes > MAX_ARRAY_SIZE:
+            raise MaxArraySizeError(
+                f"Maximum array limit {MAX_ARRAY_SIZE} reached, trying to put DataArray of {da.shape} in memory."
+            )
+
         src_width = da.rio.width
         src_height = da.rio.height
         src_bounds = list(da.rio.bounds())
@@ -374,7 +389,7 @@ class XarrayReader(BaseReader):
             src_bounds[1] = max(src_bounds[1], -85.06)
             src_bounds[3] = min(src_bounds[3], 85.06)
 
-            # South->North
+            # North->South
             if src_transform.e > 0:
                 src_bounds = [src_bounds[0], src_bounds[3], src_bounds[2], src_bounds[1]]
             # West->East
@@ -422,45 +437,33 @@ class XarrayReader(BaseReader):
         )
 
         arr = da.to_masked_array()
+        arr.mask |= arr.data == da.rio.nodata
         if out_dtype:
             arr = arr.astype(out_dtype)
-        arr.mask |= arr.data == da.rio.nodata
 
         output_bounds = da.rio._unordered_bounds()
         if output_bounds[1] > output_bounds[3] and da.rio.transform().e > 0:
             yaxis = self.input.dims.index(self.input.rio.y_dim)
             arr = numpy.flip(arr, axis=yaxis)
 
-        img = ImageData(
+        return ImageData(
             arr,
             bounds=bbox,
             crs=da.rio.crs,
             dataset_statistics=stats,
             band_names=band_names,
+            band_descriptions=band_descriptions,
+            nodata=da.rio.nodata,
         )
 
-        output_height = height or img.height
-        output_width = width or img.width
-        if max_size and not (width and height):
-            output_height, output_width = _get_width_height(
-                max_size, img.height, img.width
-            )
-
-        if output_height != img.height or output_width != img.width:
-            img = img.resize(
-                output_height, output_width, resampling_method=resampling_method
-            )
-
-        return img
-
-    def preview(
+    def preview(  # noqa: C901
         self,
         max_size: int = 1024,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
-        nodata: Optional[NoData] = None,
-        indexes: Optional[Indexes] = None,
-        dst_crs: Optional[CRS] = None,
+        height: int | None = None,
+        width: int | None = None,
+        nodata: NoData | None = None,
+        indexes: Indexes | None = None,
+        dst_crs: CRS | None = None,
         reproject_method: WarpResampling = "nearest",
         resampling_method: RIOResampling = "nearest",
         out_dtype: str | numpy.dtype | None = None,
@@ -488,7 +491,12 @@ class XarrayReader(BaseReader):
             )
             max_size = None
 
-        da, band_names = self._sel_indexes(indexes)
+        da, band_names, band_descriptions = self._sel_indexes(indexes)
+
+        if da.nbytes > MAX_ARRAY_SIZE:
+            raise MaxArraySizeError(
+                f"Maximum array limit {MAX_ARRAY_SIZE} reached, trying to put DataArray of {da.shape} in memory."
+            )
 
         if nodata is not None:
             da = da.rio.write_nodata(nodata)
@@ -574,6 +582,8 @@ class XarrayReader(BaseReader):
             crs=da.rio.crs,
             dataset_statistics=stats,
             band_names=band_names,
+            band_descriptions=band_descriptions,
+            nodata=da.rio.nodata,
         )
 
         if max_size:
@@ -596,8 +606,8 @@ class XarrayReader(BaseReader):
         lon: float,
         lat: float,
         coord_crs: CRS = WGS84_CRS,
-        nodata: Optional[NoData] = None,
-        indexes: Optional[Indexes] = None,
+        nodata: NoData | None = None,
+        indexes: Indexes | None = None,
         out_dtype: str | numpy.dtype | None = None,
         **kwargs: Any,
     ) -> PointData:
@@ -621,7 +631,7 @@ class XarrayReader(BaseReader):
         ):
             raise PointOutsideBounds("Point is outside dataset bounds")
 
-        da, band_names = self._sel_indexes(indexes)
+        da, band_names, band_descriptions = self._sel_indexes(indexes)
 
         if nodata is not None:
             da = da.rio.write_nodata(nodata)
@@ -642,21 +652,23 @@ class XarrayReader(BaseReader):
             coordinates=(lon, lat),
             crs=coord_crs,
             band_names=band_names,
+            band_descriptions=band_descriptions,
             pixel_location=(x, y),
+            nodata=da.rio.nodata,
         )
 
     def feature(
         self,
         shape: Dict,
-        dst_crs: Optional[CRS] = None,
+        dst_crs: CRS | None = None,
         shape_crs: CRS = WGS84_CRS,
         reproject_method: WarpResampling = "nearest",
         auto_expand: bool = True,
-        nodata: Optional[NoData] = None,
-        indexes: Optional[Indexes] = None,
-        max_size: Optional[int] = None,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
+        nodata: NoData | None = None,
+        indexes: Indexes | None = None,
+        max_size: int | None = None,
+        height: int | None = None,
+        width: int | None = None,
         resampling_method: RIOResampling = "nearest",
         out_dtype: str | numpy.dtype | None = None,
         **kwargs: Any,