include #883

hydromt/data_adapter/rasterdataset.py

@@ -267,19 +267,20 @@ def _slice_data(
         ds : xarray.Dataset
             The sliced RasterDataset.
         """
-        if isinstance(ds, xr.DataArray):
+        if isinstance(ds, xr.DataArray):  # xr.DataArray has no variables
             if ds.name is None:
                 # dummy name, required to create dataset
                 # renamed to variable in _single_var_as_array
                 ds.name = "data"
             ds = ds.to_dataset()
-        elif variables is not None:
-            variables = np.atleast_1d(variables).tolist()
-            if len(variables) > 1 or len(ds.data_vars) > 1:
-                mvars = [var not in ds.data_vars for var in variables]
-                if any(mvars):
-                    raise ValueError(f"RasterDataset: variables not found {mvars}")
-                ds = ds[variables]
+        elif variables is not None:  # xr.Dataset has variables
+            # cast variables to list
+            variables = list([variables]) if isinstance(variables, str) else variables
+            mvars = [var for var in variables if var not in ds.data_vars]
+            if len(mvars) > 0:
+                raise NoDataException(f"Variables {mvars} not found in data.")
+            ds = ds[variables]
+
         if time_tuple is not None:
             ds = RasterDatasetAdapter._slice_temporal_dimension(
                 ds,

hydromt/drivers/raster_xarray_driver.py

@@ -8,7 +8,7 @@
 
 import xarray as xr
 
-from hydromt._typing import Geom, StrPath, TimeRange, ZoomLevel
+from hydromt._typing import Geom, StrPath, TimeRange, Variables, ZoomLevel
 from hydromt._typing.error import NoDataStrategy
 from hydromt._utils.unused_kwargs import warn_on_unused_kwargs
 from hydromt.drivers.preprocessing import PREPROCESSORS
@@ -26,6 +26,7 @@ def read_data(
         *,
         logger: Logger,
         mask: Optional[Geom] = None,
+        variables: Optional[Variables] = None,
         time_range: Optional[TimeRange] = None,
         zoom_level: Optional[ZoomLevel] = None,
         handle_nodata: NoDataStrategy = NoDataStrategy.RAISE,
@@ -38,7 +39,12 @@ def read_data(
         """
         warn_on_unused_kwargs(
             self.__class__.__name__,
-            {"mask": mask, "time_range": time_range, "zoom_level": zoom_level},
+            {
+                "mask": mask,
+                "time_range": time_range,
+                "variables": variables,
+                "zoom_level": zoom_level,
+            },
             logger,
         )
         options = copy(self.options)

hydromt/drivers/rasterdataset_driver.py

@@ -6,7 +6,7 @@
 
 import xarray as xr
 
-from hydromt._typing import Geom, StrPath, TimeRange, ZoomLevel
+from hydromt._typing import Geom, StrPath, TimeRange, Variables, ZoomLevel
 from hydromt._typing.error import NoDataStrategy
 
 from .base_driver import BaseDriver
@@ -20,7 +20,7 @@ def read(
         uri: str,
         *,
         mask: Optional[Geom] = None,
-        variables: Optional[List[str]] = None,
+        variables: Optional[Variables] = None,
         time_range: Optional[TimeRange] = None,
         zoom_level: Optional[ZoomLevel] = None,
         logger: Optional[Logger] = None,
@@ -48,6 +48,7 @@ def read(
             uris,
             mask=mask,
             time_range=time_range,
+            variables=variables,
             zoom_level=zoom_level,
             logger=logger,
             handle_nodata=handle_nodata,
@@ -59,6 +60,7 @@ def read_data(
         uris: List[str],
         *,
         mask: Optional[Geom] = None,
+        variables: Optional[Variables] = None,
         time_range: Optional[TimeRange] = None,
         zoom_level: Optional[ZoomLevel] = None,
         logger: Logger,

hydromt/drivers/rasterio_driver.py

@@ -12,8 +12,9 @@
 import xarray as xr
 from pyproj import CRS
 
-from hydromt._typing import Geom, StrPath, TimeRange, ZoomLevel
+from hydromt._typing import Geom, StrPath, TimeRange, Variables, ZoomLevel
 from hydromt._typing.error import NoDataStrategy
+from hydromt._utils.unused_kwargs import warn_on_unused_kwargs
 from hydromt._utils.uris import strip_scheme
 from hydromt.config import SETTINGS
 from hydromt.data_adapter.caching import cache_vrt_tiles
@@ -33,12 +34,16 @@ def read_data(
         *,
         mask: Optional[Geom] = None,
         time_range: Optional[TimeRange] = None,
+        variables: Optional[Variables] = None,
         zoom_level: Optional[ZoomLevel] = None,
         logger: Logger,
         handle_nodata: NoDataStrategy = NoDataStrategy.RAISE,
     ) -> xr.Dataset:
         """Read data using rasterio."""
         # build up kwargs for open_raster
+        warn_on_unused_kwargs(
+            self.__class__.__name__, {"time_range": time_range}, logger=logger
+        )
         kwargs: Dict[str, Any] = {}
 
         # get source-specific options
@@ -73,6 +78,9 @@ def read_data(
                 # NOTE: overview levels start at zoom_level 1, see _get_zoom_levels_and_crs
                 kwargs.update(overview_level=zoom_level - 1)
         ds = open_mfraster(uris, logger=logger, **kwargs)
+        # rename ds with single band if single variable is requested
+        if variables is not None and len(variables) == 1 and len(ds.data_vars) == 1:
+            ds = ds.rename({list(ds.data_vars.keys())[0]: list(variables)[0]})
         return ds
 
     def write(self, path: StrPath, ds: xr.Dataset, **kwargs) -> None:

tests/gis/test_raster.py

@@ -14,8 +14,8 @@
 from affine import Affine
 from shapely.geometry import LineString, Point, box
 
+from hydromt.drivers.rasterio_driver import open_raster
 from hydromt.gis import raster, utils
-from hydromt.io import open_raster
 
 # origin, rotation, res, shape, internal_bounds
 # NOTE a rotated grid with a negative dx is not supported