read image element from base store

Author: berombau

src/spatialdata/_io/io_raster.py

@@ -1,9 +1,9 @@
-from pathlib import Path
 from typing import Any, Literal
 
 import dask.array as da
 import numpy as np
 import zarr
+import zarr.storage
 from ome_zarr.format import Format
 from ome_zarr.io import ZarrLocation
 from ome_zarr.reader import Label, Multiscales, Node, Reader
@@ -15,16 +15,8 @@
 from ome_zarr.writer import write_multiscale_labels as write_multiscale_labels_ngff
 from xarray import DataArray, Dataset, DataTree
 
-from spatialdata._io._utils import (
-    _get_transformations_from_ngff_dict,
-    overwrite_coordinate_transformations_raster,
-)
-from spatialdata._io.format import (
-    CurrentRasterFormat,
-    RasterFormats,
-    RasterFormatV01,
-    _parse_version,
-)
+from spatialdata._io._utils import _get_transformations_from_ngff_dict, overwrite_coordinate_transformations_raster
+from spatialdata._io.format import CurrentRasterFormat, RasterFormats, RasterFormatV01, _parse_version
 from spatialdata._utils import get_pyramid_levels
 from spatialdata.models._utils import get_channel_names
 from spatialdata.models.models import ATTRS_KEY
@@ -36,16 +28,16 @@
 )
 
 
-def _read_multiscale(store: str | Path, raster_type: Literal["image", "labels"]) -> DataArray | DataTree:
-    assert isinstance(store, str | Path)
+def _read_multiscale(store: zarr.storage.BaseStore, raster_type: Literal["image", "labels"]) -> DataArray | DataTree:
+    assert isinstance(store, zarr.storage.BaseStore)
     assert raster_type in ["image", "labels"]
 
-    f = zarr.open(store, mode="r")
-    version = _parse_version(f, expect_attrs_key=True)
+    group = zarr.group(store=store)
+    version = _parse_version(group, expect_attrs_key=True)
     # old spatialdata datasets don't have format metadata for raster elements; this line ensure backwards compatibility,
     # interpreting the lack of such information as the presence of the format v01
     format = RasterFormatV01() if version is None else RasterFormats[version]
-    f.store.close()
+    store.close()
 
     nodes: list[Node] = []
     image_loc = ZarrLocation(store)