add decorator skipping nodes without dimensions

Author: melonora

src/spatialdata/_utils.py

@@ -311,3 +311,22 @@ def _error_message_add_element() -> None:
         "write_labels(), write_points(), write_shapes() and write_table(). We are going to make these calls more "
         "ergonomic in a follow up PR."
     )
+
+
+def skip_non_dimension_nodes(func: Callable[[Dataset], Dataset]) -> Callable[[Dataset], Dataset]:
+    """Skip nodes in Datatree that do not contain dimensions.
+
+    This function implements the workaround of https://github.com/pydata/xarray/issues/9693. In particular,
+    we need this because of our DataTree representing multiscale image having a root node that does not have
+    dimensions. Several functions need to be mapped over the datasets in the datatree that depend on having
+    dimensions, e.g. a transpose.
+    """
+
+    @functools.wraps(func)
+    def _func(ds: Dataset, *args: Any, **kwargs: Any) -> Dataset:
+        # check if dimensions are present otherwise return verbatim
+        if len(ds.dims) == 0:
+            return ds
+        return func(ds, *args, **kwargs)
+
+    return _func

tests/utils/test_element_utils.py

@@ -5,7 +5,7 @@
 import xarray
 from xarray import DataArray, DataTree
 
-from spatialdata._utils import unpad_raster
+from spatialdata._utils import skip_non_dimension_nodes, unpad_raster
 from spatialdata.models import get_model
 from spatialdata.transformations import Affine
 
@@ -64,3 +64,19 @@ def test_unpad_raster(images, labels) -> None:
                 raise e
         else:
             raise ValueError(f"Unknown type: {type(raster)}")
+
+
+def test_skip_nodes(images):
+    multiscale_img = images["image2d_multiscale"]
+
+    @skip_non_dimension_nodes
+    def transpose(ds, *args, **kwargs):
+        return ds.transpose(*args, **kwargs)
+
+    for scale in list(multiscale_img.keys()):
+        assert multiscale_img[scale]["image"].dims == ("c", "y", "x")
+
+    # applying this function without skipping the root node would fail as the root node does not have dimensions.
+    result = images["image2d_multiscale"].map_over_datasets(transpose, "y", "x", "c")
+    for scale in list(result.keys()):
+        assert result[scale]["image"].dims == ("y", "x", "c")