add _open_zarr_store

Author: berombau

src/spatialdata/_core/spatialdata.py

@@ -31,10 +31,7 @@
 )
 from spatialdata._logging import logger
 from spatialdata._types import ArrayLike, Raster_T
-from spatialdata._utils import (
-    _deprecation_alias,
-    _error_message_add_element,
-)
+from spatialdata._utils import _deprecation_alias, _error_message_add_element
 from spatialdata.models import (
     Image2DModel,
     Image3DModel,
@@ -600,7 +597,7 @@ def path(self, value: Path | None) -> None:
             )
 
     def _get_groups_for_element(
-        self, zarr_path: Path, element_type: str, element_name: str
+        self, zarr_path: UPath, element_type: str, element_name: str
     ) -> tuple[zarr.Group, zarr.Group, zarr.Group]:
         """
         Get the Zarr groups for the root, element_type and element for a specific element.
@@ -620,9 +617,9 @@ def _get_groups_for_element(
         -------
         either the existing Zarr subgroup or a new one.
         """
-        if not isinstance(zarr_path, Path):
-            raise ValueError("zarr_path should be a Path object")
-        store = parse_url(zarr_path, mode="r+").store
+        from spatialdata._io._utils import _open_zarr_store
+
+        store = _open_zarr_store(zarr_path, mode="r+")
         root = zarr.group(store=store)
         if element_type not in ["images", "labels", "points", "polygons", "shapes", "tables"]:
             raise ValueError(f"Unknown element type {element_type}")
@@ -1375,7 +1372,7 @@ def delete_element_from_disk(self, element_name: str | list[str]) -> None:
                 self.delete_element_from_disk(name)
             return
 
-        from spatialdata._io._utils import _backed_elements_contained_in_path
+        from spatialdata._io._utils import _backed_elements_contained_in_path, _open_zarr_store
 
         if self.path is None:
             raise ValueError("The SpatialData object is not backed by a Zarr store.")
@@ -1416,7 +1413,7 @@ def delete_element_from_disk(self, element_name: str | list[str]) -> None:
             )
 
         # delete the element
-        store = parse_url(self.path, mode="r+").store
+        store = _open_zarr_store(self.path)
         root = zarr.group(store=store)
         root[element_type].pop(element_name)
         store.close()
@@ -1437,7 +1434,9 @@ def _check_element_not_on_disk_with_different_type(self, element_type: str, elem
                 )
 
     def write_consolidated_metadata(self) -> None:
-        store = parse_url(self.path, mode="r+").store
+        from spatialdata._io._utils import _open_zarr_store
+
+        store = _open_zarr_store(self.path)
         # consolidate metadata to more easily support remote reading bug in zarr. In reality, 'zmetadata' is written
         # instead of '.zmetadata' see discussion https://github.com/zarr-developers/zarr-python/issues/1121
         zarr.consolidate_metadata(store, metadata_key=".zmetadata")
@@ -1574,15 +1573,11 @@ def write_transformations(self, element_name: str | None = None) -> None:
         )
         axes = get_axes_names(element)
         if isinstance(element, DataArray | DataTree):
-            from spatialdata._io._utils import (
-                overwrite_coordinate_transformations_raster,
-            )
+            from spatialdata._io._utils import overwrite_coordinate_transformations_raster
 
             overwrite_coordinate_transformations_raster(group=element_group, axes=axes, transformations=transformations)
         elif isinstance(element, DaskDataFrame | GeoDataFrame | AnnData):
-            from spatialdata._io._utils import (
-                overwrite_coordinate_transformations_non_raster,
-            )
+            from spatialdata._io._utils import overwrite_coordinate_transformations_non_raster
 
             overwrite_coordinate_transformations_non_raster(
                 group=element_group, axes=axes, transformations=transformations

src/spatialdata/_io/_utils.py

@@ -17,7 +17,10 @@
 from dask.array import Array as DaskArray
 from dask.dataframe import DataFrame as DaskDataFrame
 from geopandas import GeoDataFrame
+from upath import UPath
+from upath.implementations.local import PosixUPath, WindowsUPath
 from xarray import DataArray, DataTree
+from zarr.storage import FSStore
 
 from spatialdata._core.spatialdata import SpatialData
 from spatialdata._utils import get_pyramid_levels
@@ -28,10 +31,7 @@
     _validate_mapping_to_coordinate_system_type,
 )
 from spatialdata.transformations.ngff.ngff_transformations import NgffBaseTransformation
-from spatialdata.transformations.transformations import (
-    BaseTransformation,
-    _get_current_output_axes,
-)
+from spatialdata.transformations.transformations import BaseTransformation, _get_current_output_axes
 
 
 # suppress logger debug from ome_zarr with context manager
@@ -383,3 +383,12 @@ def save_transformations(sdata: SpatialData) -> None:
         stacklevel=2,
     )
     sdata.write_transformations()
+
+
+def _open_zarr_store(path: str | UPath, **kwargs) -> zarr.storage.BaseStore:
+    if isinstance(path, str | Path):
+        path = UPath(path)
+    if isinstance(path, PosixUPath | WindowsUPath):
+        return zarr.storage.DirectoryStore(path.path)
+    else:
+        return FSStore(path.path, fs=path.fs, **kwargs)

tests/io/test_remote_mock.py

@@ -263,6 +263,9 @@ def test_local_sdata_remote_image(self, upath: UPath, images: SpatialData) -> No
             local_sdata = SpatialData.read(sdata_path)  # noqa: F841
             remote_path = upath / "full_sdata.zarr"  # noqa: F841
 
+            remote_sdata = SpatialData.read(remote_path)
+            assert_spatial_data_objects_are_identical(local_sdata, remote_sdata)
+
             # TODO: read a single remote image from the S3 data and add it to the local SpatialData object
             # for a in remote_path.glob('**/*'):
             #     print(a)