[scverse hackathon] cloud future

pyproject.toml

@@ -26,20 +26,22 @@ dependencies = [
  "click",
  "dask-image",
  "dask>=2024.4.1,<=2024.11.2",
- "fsspec",
+ "fsspec[s3,http]",
  "geopandas>=0.14",
  "multiscale_spatial_image>=2.0.2",
  "networkx",
  "numba>=0.55.0",
  "numpy",
  "ome_zarr>=0.8.4",
+ "universal_pathlib>=0.2.6",
  "pandas",
  "pooch",
  "pyarrow",
  "rich",
  "setuptools",
  "shapely>=2.0.1",
- "spatial_image>=1.1.0",
+ "spatial_image>=1.2.1",
+ "xarray-dataclasses>=1.9.1",
  "scikit-image",
  "scipy",
  "typing_extensions>=4.8.0",
@@ -58,6 +60,7 @@ test = [
  "pytest-cov",
  "pytest-mock",
  "torch",
+ "moto[s3,server]"
 ]
 docs = [
  "sphinx>=4.5",

src/spatialdata/_core/spatialdata.py

@@ -30,11 +30,8 @@
  validate_table_attr_keys,
 )
 from spatialdata._logging import logger
-from spatialdata._types import ArrayLike, Raster_T
-from spatialdata._utils import (
- _deprecation_alias,
- _error_message_add_element,
-)
+from spatialdata._types import ArrayLike, Raster_T, StoreLike
+from spatialdata._utils import _deprecation_alias, _error_message_add_element
 from spatialdata.models import (
  Image2DModel,
  Image3DModel,
@@ -601,7 +598,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: StoreLike, 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.
@@ -621,9 +618,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}")
@@ -1068,9 +1065,12 @@ def elements_paths_on_disk(self) -> list[str]:
  -------
  A list of paths of the elements saved in the Zarr store.
  """
+ from spatialdata._io._utils import _open_zarr_store
+
  if self.path is None:
  raise ValueError("The SpatialData object is not backed by a Zarr store.")
- store = parse_url(self.path, mode="r").store
+
+ store = _open_zarr_store(self.path)
  root = zarr.group(store=store)
  elements_in_zarr = []
 
@@ -1175,7 +1175,7 @@ def _validate_all_elements(self) -> None:
 
  def write(
  self,
- file_path: str | Path,
+ file_path: StoreLike,
  overwrite: bool = False,
  consolidate_metadata: bool = True,
  format: SpatialDataFormat | list[SpatialDataFormat] | None = None,
@@ -1205,12 +1205,16 @@ def write(
  :class:`~spatialdata._io.format.CurrentRasterFormat`, :class:`~spatialdata._io.format.CurrentShapesFormat`,
  :class:`~spatialdata._io.format.CurrentPointsFormat`, :class:`~spatialdata._io.format.CurrentTablesFormat`.
  """
+ from spatialdata._io._utils import _open_zarr_store
+
  if isinstance(file_path, str):
  file_path = Path(file_path)
- self._validate_can_safely_write_to_path(file_path, overwrite=overwrite)
- self._validate_all_elements()
+ if isinstance(file_path, Path):
+ # TODO: also validate remote paths
+ self._validate_can_safely_write_to_path(file_path, overwrite=overwrite)
+ self._validate_all_elements()
 
- store = parse_url(file_path, mode="w").store
+ store = _open_zarr_store(file_path, mode="w")
  zarr_group = zarr.group(store=store, overwrite=overwrite)
  self.write_attrs(zarr_group=zarr_group)
  store.close()
@@ -1236,20 +1240,21 @@ def write(
  def _write_element(
  self,
  element: SpatialElement | AnnData,
- zarr_container_path: Path,
+ zarr_container_path: StoreLike,
  element_type: str,
  element_name: str,
  overwrite: bool,
  format: SpatialDataFormat | list[SpatialDataFormat] | None = None,
  ) -> None:
- if not isinstance(zarr_container_path, Path):
+ if not isinstance(zarr_container_path, StoreLike):
  raise ValueError(
  f"zarr_container_path must be a Path object, type(zarr_container_path) = {type(zarr_container_path)}."
  )
- file_path_of_element = zarr_container_path / element_type / element_name
- self._validate_can_safely_write_to_path(
- file_path=file_path_of_element, overwrite=overwrite, saving_an_element=True
- )
+ if isinstance(zarr_container_path, Path):
+ file_path_of_element = zarr_container_path / element_type / element_name
+ self._validate_can_safely_write_to_path(
+ file_path=file_path_of_element, overwrite=overwrite, saving_an_element=True
+ )
 
  root_group, element_type_group, _ = self._get_groups_for_element(
  zarr_path=zarr_container_path, element_type=element_type, element_name=element_name
@@ -1376,7 +1381,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.")
@@ -1417,7 +1422,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()
@@ -1438,15 +1443,24 @@ 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
- # 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")
+ from spatialdata._io._utils import _open_zarr_store
+
+ store = _open_zarr_store(self.path)
+ # Note that the store can be local (which does not have the zmetadata bug)
+ # or a remote FSStore (which has the bug).
+ # Consolidate metadata to more easily support remote reading bug in zarr.
+ # We write 'zmetadata' instead of the standard '.zmetadata' to avoid the FSStore bug.
+ # See discussion https://github.com/zarr-developers/zarr-python/issues/1121
+ zarr.consolidate_metadata(store, metadata_key="zmetadata")
  store.close()
 
  def has_consolidated_metadata(self) -> bool:
+ from spatialdata._io._utils import _open_zarr_store
+
  return_value = False
- store = parse_url(self.path, mode="r").store
+ store = _open_zarr_store(self.path)
+ # Note that the store can be local (which does not have the zmetadata bug)
+ # or a remote FSStore (which has the bug).
  if "zmetadata" in store:
  return_value = True
  store.close()
@@ -1575,15 +1589,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

@@ -15,14 +15,18 @@
 from pathlib import Path
 from typing import Any, Literal
 
-import zarr
+import zarr.storage
 from anndata import AnnData
 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._types import StoreLike
 from spatialdata._utils import get_pyramid_levels
 from spatialdata.models._utils import (
  MappingToCoordinateSystem_t,
@@ -31,10 +35,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
@@ -388,6 +389,48 @@ def save_transformations(sdata: SpatialData) -> None:
  sdata.write_transformations()
 
 
+def _open_zarr_store(path: StoreLike, **kwargs: Any) -> zarr.storage.BaseStore:
+ # TODO: ensure kwargs like mode are enforced everywhere and passed correctly to the store
+ if isinstance(path, str | Path):
+ # if the input is str or Path, map it to UPath
+ path = UPath(path)
+ if isinstance(path, PosixUPath | WindowsUPath):
+ # if the input is a local path, use DirectoryStore
+ return zarr.storage.DirectoryStore(path.path, dimension_separator="/")
+ if isinstance(path, zarr.Group):
+ # if the input is a zarr.Group, wrap it with a store
+ if isinstance(path.store, zarr.storage.DirectoryStore):
+ # create a simple FSStore if the store is a DirectoryStore with just the path
+ return FSStore(os.path.join(path.store.path, path.path), **kwargs)
+ if isinstance(path.store, FSStore):
+ # if the store within the zarr.Group is an FSStore, return it
+ # but extend the path of the store with that of the zarr.Group
+ return FSStore(path.store.path + "/" + path.path, fs=path.store.fs, **kwargs)
+ if isinstance(path.store, zarr.storage.ConsolidatedMetadataStore):
+ # if the store is a ConsolidatedMetadataStore, just return the underlying FSSpec store
+ return path.store.store
+ raise ValueError(f"Unsupported store type or zarr.Group: {type(path.store)}")
+ if isinstance(path, zarr.storage.StoreLike):
+ # if the input already a store, wrap it in an FSStore
+ return FSStore(path, **kwargs)
+ if isinstance(path, UPath):
+ # if input is a remote UPath, map it to an FSStore
+ return FSStore(path.path, fs=path.fs, **kwargs)
+ raise TypeError(f"Unsupported type: {type(path)}")
+
+
+def _create_upath(path: StoreLike) -> UPath | None:
+ # try to create a UPath from the input
+ if isinstance(path, zarr.storage.ConsolidatedMetadataStore):
+ path = path.store # get the fsstore from the consolidated store
+ if isinstance(path, FSStore):
+ protocol = path.fs.protocol if isinstance(path.fs.protocol, str) else path.fs.protocol[0]
+ return UPath(path.path, protocol=protocol, **path.fs.storage_options)
+ if isinstance(path, zarr.storage.DirectoryStore):
+ return UPath(path.path)
+ return None
+
+
 class BadFileHandleMethod(Enum):
  ERROR = "error"
  WARN = "warn"

src/spatialdata/_io/format.py

@@ -257,7 +257,7 @@ def validate_table(
 def format_implementations() -> Iterator[Format]:
  """Return an instance of each format implementation, newest to oldest."""
  yield RasterFormatV02()
-# yield RasterFormatV01() # same format string as FormatV04
+ yield RasterFormatV01() # same format string as FormatV04
  yield FormatV04()
  yield FormatV03()
  yield FormatV02()

src/spatialdata/_io/io_points.py

@@ -14,27 +14,19 @@
 )
 from spatialdata._io.format import CurrentPointsFormat, PointsFormats, _parse_version
 from spatialdata.models import get_axes_names
-from spatialdata.transformations._utils import (
- _get_transformations,
- _set_transformations,
-)
+from spatialdata.transformations._utils import _get_transformations, _set_transformations
 
 
 def _read_points(
  store: str | Path | MutableMapping | zarr.Group, # type: ignore[type-arg]
 ) -> DaskDataFrame:
  """Read points from a zarr store."""
- assert isinstance(store, str | Path)
- f = zarr.open(store, mode="r")
-
+ f = zarr.open(store, mode="r") if isinstance(store, str | Path | MutableMapping) else store
  version = _parse_version(f, expect_attrs_key=True)
  assert version is not None
  format = PointsFormats[version]
 
- path = os.path.join(f._store.path, f.path, "points.parquet")
- # cache on remote file needed for parquet reader to work
- # TODO: allow reading in the metadata without caching all the data
- points = read_parquet("simplecache::" + path if path.startswith("http") else path)
+ points = read_parquet(f.store.path, filesystem=getattr(f.store, "fs", None))
  assert isinstance(points, DaskDataFrame)
 
  transformations = _get_transformations_from_ngff_dict(f.attrs.asdict()["coordinateTransformations"])
@@ -57,7 +49,8 @@ def write_points(
  t = _get_transformations(points)
 
  points_groups = group.require_group(name)
- path = Path(points_groups._store.path) / points_groups.path / "points.parquet"
+ store = points_groups._store
+ new_path = os.path.join(store.path, points_groups.path, "points.parquet")
 
  # The following code iterates through all columns in the 'points' DataFrame. If the column's datatype is
  # 'category', it checks whether the categories of this column are known. If not, it explicitly converts the
@@ -70,7 +63,7 @@ def write_points(
  c = c.cat.as_known()
  points[column_name] = c
 
- points.to_parquet(path)
+ points.to_parquet(new_path, filesystem=getattr(store, "fs", None))
 
  attrs = format.attrs_to_dict(points.attrs)
  attrs["version"] = format.spatialdata_format_version

pyproject.toml

@@ -26,20 +26,22 @@ dependencies = [
  "click",
  "dask-image",
  "dask>=2024.4.1,<=2024.11.2",
- "fsspec",
+ "fsspec[s3,http]",
  "geopandas>=0.14",
  "multiscale_spatial_image>=2.0.2",
  "networkx",
  "numba>=0.55.0",
  "numpy",
  "ome_zarr>=0.8.4",
+ "universal_pathlib>=0.2.6",
  "pandas",
  "pooch",
  "pyarrow",
  "rich",
  "setuptools",
  "shapely>=2.0.1",
- "spatial_image>=1.1.0",
+ "spatial_image>=1.2.1",
+ "xarray-dataclasses>=1.9.1",
  "scikit-image",
  "scipy",
  "typing_extensions>=4.8.0",
@@ -58,6 +60,7 @@ test = [
  "pytest-cov",
  "pytest-mock",
  "torch",
+ "moto[s3,server]"
 ]
 docs = [
  "sphinx>=4.5",

src/spatialdata/_core/spatialdata.py

@@ -30,11 +30,8 @@
  validate_table_attr_keys,
 )
 from spatialdata._logging import logger
-from spatialdata._types import ArrayLike, Raster_T
-from spatialdata._utils import (
- _deprecation_alias,
- _error_message_add_element,
-)
+from spatialdata._types import ArrayLike, Raster_T, StoreLike
+from spatialdata._utils import _deprecation_alias, _error_message_add_element
 from spatialdata.models import (
  Image2DModel,
  Image3DModel,
@@ -601,7 +598,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: StoreLike, 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.
@@ -621,9 +618,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}")
@@ -1068,9 +1065,12 @@ def elements_paths_on_disk(self) -> list[str]:
  -------
  A list of paths of the elements saved in the Zarr store.
  """
+ from spatialdata._io._utils import _open_zarr_store
+
  if self.path is None:
  raise ValueError("The SpatialData object is not backed by a Zarr store.")
- store = parse_url(self.path, mode="r").store
+
+ store = _open_zarr_store(self.path)
  root = zarr.group(store=store)
  elements_in_zarr = []
 
@@ -1175,7 +1175,7 @@ def _validate_all_elements(self) -> None:
 
  def write(
  self,
- file_path: str | Path,
+ file_path: StoreLike,
  overwrite: bool = False,
  consolidate_metadata: bool = True,
  format: SpatialDataFormat | list[SpatialDataFormat] | None = None,
@@ -1205,12 +1205,16 @@ def write(
  :class:`~spatialdata._io.format.CurrentRasterFormat`, :class:`~spatialdata._io.format.CurrentShapesFormat`,
  :class:`~spatialdata._io.format.CurrentPointsFormat`, :class:`~spatialdata._io.format.CurrentTablesFormat`.
  """
+ from spatialdata._io._utils import _open_zarr_store
+
  if isinstance(file_path, str):
  file_path = Path(file_path)
- self._validate_can_safely_write_to_path(file_path, overwrite=overwrite)
- self._validate_all_elements()
+ if isinstance(file_path, Path):
+ # TODO: also validate remote paths
+ self._validate_can_safely_write_to_path(file_path, overwrite=overwrite)
+ self._validate_all_elements()
 
- store = parse_url(file_path, mode="w").store
+ store = _open_zarr_store(file_path, mode="w")
  zarr_group = zarr.group(store=store, overwrite=overwrite)
  self.write_attrs(zarr_group=zarr_group)
  store.close()
@@ -1236,20 +1240,21 @@ def write(
  def _write_element(
  self,
  element: SpatialElement | AnnData,
- zarr_container_path: Path,
+ zarr_container_path: StoreLike,
  element_type: str,
  element_name: str,
  overwrite: bool,
  format: SpatialDataFormat | list[SpatialDataFormat] | None = None,
  ) -> None:
- if not isinstance(zarr_container_path, Path):
+ if not isinstance(zarr_container_path, StoreLike):
  raise ValueError(
  f"zarr_container_path must be a Path object, type(zarr_container_path) = {type(zarr_container_path)}."
  )
- file_path_of_element = zarr_container_path / element_type / element_name
- self._validate_can_safely_write_to_path(
- file_path=file_path_of_element, overwrite=overwrite, saving_an_element=True
- )
+ if isinstance(zarr_container_path, Path):
+ file_path_of_element = zarr_container_path / element_type / element_name
+ self._validate_can_safely_write_to_path(
+ file_path=file_path_of_element, overwrite=overwrite, saving_an_element=True
+ )
 
  root_group, element_type_group, _ = self._get_groups_for_element(
  zarr_path=zarr_container_path, element_type=element_type, element_name=element_name
@@ -1376,7 +1381,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.")
@@ -1417,7 +1422,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()
@@ -1438,15 +1443,24 @@ 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
- # 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")
+ from spatialdata._io._utils import _open_zarr_store
+
+ store = _open_zarr_store(self.path)
+ # Note that the store can be local (which does not have the zmetadata bug)
+ # or a remote FSStore (which has the bug).
+ # Consolidate metadata to more easily support remote reading bug in zarr.
+ # We write 'zmetadata' instead of the standard '.zmetadata' to avoid the FSStore bug.
+ # See discussion https://github.com/zarr-developers/zarr-python/issues/1121
+ zarr.consolidate_metadata(store, metadata_key="zmetadata")
  store.close()
 
  def has_consolidated_metadata(self) -> bool:
+ from spatialdata._io._utils import _open_zarr_store
+
  return_value = False
- store = parse_url(self.path, mode="r").store
+ store = _open_zarr_store(self.path)
+ # Note that the store can be local (which does not have the zmetadata bug)
+ # or a remote FSStore (which has the bug).
  if "zmetadata" in store:
  return_value = True
  store.close()
@@ -1575,15 +1589,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

@@ -15,14 +15,18 @@
 from pathlib import Path
 from typing import Any, Literal
 
-import zarr
+import zarr.storage
 from anndata import AnnData
 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._types import StoreLike
 from spatialdata._utils import get_pyramid_levels
 from spatialdata.models._utils import (
  MappingToCoordinateSystem_t,
@@ -31,10 +35,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
@@ -388,6 +389,48 @@ def save_transformations(sdata: SpatialData) -> None:
  sdata.write_transformations()
 
 
+def _open_zarr_store(path: StoreLike, **kwargs: Any) -> zarr.storage.BaseStore:
+ # TODO: ensure kwargs like mode are enforced everywhere and passed correctly to the store
+ if isinstance(path, str | Path):
+ # if the input is str or Path, map it to UPath
+ path = UPath(path)
+ if isinstance(path, PosixUPath | WindowsUPath):
+ # if the input is a local path, use DirectoryStore
+ return zarr.storage.DirectoryStore(path.path, dimension_separator="/")
+ if isinstance(path, zarr.Group):
+ # if the input is a zarr.Group, wrap it with a store
+ if isinstance(path.store, zarr.storage.DirectoryStore):
+ # create a simple FSStore if the store is a DirectoryStore with just the path
+ return FSStore(os.path.join(path.store.path, path.path), **kwargs)
+ if isinstance(path.store, FSStore):
+ # if the store within the zarr.Group is an FSStore, return it
+ # but extend the path of the store with that of the zarr.Group
+ return FSStore(path.store.path + "/" + path.path, fs=path.store.fs, **kwargs)
+ if isinstance(path.store, zarr.storage.ConsolidatedMetadataStore):
+ # if the store is a ConsolidatedMetadataStore, just return the underlying FSSpec store
+ return path.store.store
+ raise ValueError(f"Unsupported store type or zarr.Group: {type(path.store)}")
+ if isinstance(path, zarr.storage.StoreLike):
+ # if the input already a store, wrap it in an FSStore
+ return FSStore(path, **kwargs)
+ if isinstance(path, UPath):
+ # if input is a remote UPath, map it to an FSStore
+ return FSStore(path.path, fs=path.fs, **kwargs)
+ raise TypeError(f"Unsupported type: {type(path)}")
+
+
+def _create_upath(path: StoreLike) -> UPath | None:
+ # try to create a UPath from the input
+ if isinstance(path, zarr.storage.ConsolidatedMetadataStore):
+ path = path.store # get the fsstore from the consolidated store
+ if isinstance(path, FSStore):
+ protocol = path.fs.protocol if isinstance(path.fs.protocol, str) else path.fs.protocol[0]
+ return UPath(path.path, protocol=protocol, **path.fs.storage_options)
+ if isinstance(path, zarr.storage.DirectoryStore):
+ return UPath(path.path)
+ return None
+
+
 class BadFileHandleMethod(Enum):
  ERROR = "error"
  WARN = "warn"

src/spatialdata/_io/format.py

@@ -257,7 +257,7 @@ def validate_table(
 def format_implementations() -> Iterator[Format]:
  """Return an instance of each format implementation, newest to oldest."""
  yield RasterFormatV02()
-# yield RasterFormatV01() # same format string as FormatV04
+ yield RasterFormatV01() # same format string as FormatV04
  yield FormatV04()
  yield FormatV03()
  yield FormatV02()

src/spatialdata/_io/io_points.py

@@ -14,27 +14,19 @@
 )
 from spatialdata._io.format import CurrentPointsFormat, PointsFormats, _parse_version
 from spatialdata.models import get_axes_names
-from spatialdata.transformations._utils import (
- _get_transformations,
- _set_transformations,
-)
+from spatialdata.transformations._utils import _get_transformations, _set_transformations
 
 
 def _read_points(
  store: str | Path | MutableMapping | zarr.Group, # type: ignore[type-arg]
 ) -> DaskDataFrame:
  """Read points from a zarr store."""
- assert isinstance(store, str | Path)
- f = zarr.open(store, mode="r")
-
+ f = zarr.open(store, mode="r") if isinstance(store, str | Path | MutableMapping) else store
  version = _parse_version(f, expect_attrs_key=True)
  assert version is not None
  format = PointsFormats[version]
 
- path = os.path.join(f._store.path, f.path, "points.parquet")
- # cache on remote file needed for parquet reader to work
- # TODO: allow reading in the metadata without caching all the data
- points = read_parquet("simplecache::" + path if path.startswith("http") else path)
+ points = read_parquet(f.store.path, filesystem=getattr(f.store, "fs", None))
  assert isinstance(points, DaskDataFrame)
 
  transformations = _get_transformations_from_ngff_dict(f.attrs.asdict()["coordinateTransformations"])
@@ -57,7 +49,8 @@ def write_points(
  t = _get_transformations(points)
 
  points_groups = group.require_group(name)
- path = Path(points_groups._store.path) / points_groups.path / "points.parquet"
+ store = points_groups._store
+ new_path = os.path.join(store.path, points_groups.path, "points.parquet")
 
  # The following code iterates through all columns in the 'points' DataFrame. If the column's datatype is
  # 'category', it checks whether the categories of this column are known. If not, it explicitly converts the
@@ -70,7 +63,7 @@ def write_points(
  c = c.cat.as_known()
  points[column_name] = c
 
- points.to_parquet(path)
+ points.to_parquet(new_path, filesystem=getattr(store, "fs", None))
 
  attrs = format.attrs_to_dict(points.attrs)
  attrs["version"] = format.spatialdata_format_version

pyproject.toml

@@ -26,20 +26,22 @@ dependencies = [
  "click",
  "dask-image",
  "dask>=2024.4.1,<=2024.11.2",
- "fsspec",
+ "fsspec[s3,http]",
  "geopandas>=0.14",
  "multiscale_spatial_image>=2.0.2",
  "networkx",
  "numba>=0.55.0",
  "numpy",
  "ome_zarr>=0.8.4",
+ "universal_pathlib>=0.2.6",
  "pandas",
  "pooch",
  "pyarrow",
  "rich",
  "setuptools",
  "shapely>=2.0.1",
- "spatial_image>=1.1.0",
+ "spatial_image>=1.2.1",
+ "xarray-dataclasses>=1.9.1",
  "scikit-image",
  "scipy",
  "typing_extensions>=4.8.0",
@@ -58,6 +60,7 @@ test = [
  "pytest-cov",
  "pytest-mock",
  "torch",
+ "moto[s3,server]"
 ]
 docs = [
  "sphinx>=4.5",

src/spatialdata/_core/spatialdata.py

@@ -30,11 +30,8 @@
  validate_table_attr_keys,
 )
 from spatialdata._logging import logger
-from spatialdata._types import ArrayLike, Raster_T
-from spatialdata._utils import (
- _deprecation_alias,
- _error_message_add_element,
-)
+from spatialdata._types import ArrayLike, Raster_T, StoreLike
+from spatialdata._utils import _deprecation_alias, _error_message_add_element
 from spatialdata.models import (
  Image2DModel,
  Image3DModel,
@@ -601,7 +598,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: StoreLike, 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.
@@ -621,9 +618,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}")
@@ -1068,9 +1065,12 @@ def elements_paths_on_disk(self) -> list[str]:
  -------
  A list of paths of the elements saved in the Zarr store.
  """
+ from spatialdata._io._utils import _open_zarr_store
+
  if self.path is None:
  raise ValueError("The SpatialData object is not backed by a Zarr store.")
- store = parse_url(self.path, mode="r").store
+
+ store = _open_zarr_store(self.path)
  root = zarr.group(store=store)
  elements_in_zarr = []
 
@@ -1175,7 +1175,7 @@ def _validate_all_elements(self) -> None:
 
  def write(
  self,
- file_path: str | Path,
+ file_path: StoreLike,
  overwrite: bool = False,
  consolidate_metadata: bool = True,
  format: SpatialDataFormat | list[SpatialDataFormat] | None = None,
@@ -1205,12 +1205,16 @@ def write(
  :class:`~spatialdata._io.format.CurrentRasterFormat`, :class:`~spatialdata._io.format.CurrentShapesFormat`,
  :class:`~spatialdata._io.format.CurrentPointsFormat`, :class:`~spatialdata._io.format.CurrentTablesFormat`.
  """
+ from spatialdata._io._utils import _open_zarr_store
+
  if isinstance(file_path, str):
  file_path = Path(file_path)
- self._validate_can_safely_write_to_path(file_path, overwrite=overwrite)
- self._validate_all_elements()
+ if isinstance(file_path, Path):
+ # TODO: also validate remote paths
+ self._validate_can_safely_write_to_path(file_path, overwrite=overwrite)
+ self._validate_all_elements()
 
- store = parse_url(file_path, mode="w").store
+ store = _open_zarr_store(file_path, mode="w")
  zarr_group = zarr.group(store=store, overwrite=overwrite)
  self.write_attrs(zarr_group=zarr_group)
  store.close()
@@ -1236,20 +1240,21 @@ def write(
  def _write_element(
  self,
  element: SpatialElement | AnnData,
- zarr_container_path: Path,
+ zarr_container_path: StoreLike,
  element_type: str,
  element_name: str,
  overwrite: bool,
  format: SpatialDataFormat | list[SpatialDataFormat] | None = None,
  ) -> None:
- if not isinstance(zarr_container_path, Path):
+ if not isinstance(zarr_container_path, StoreLike):
  raise ValueError(
  f"zarr_container_path must be a Path object, type(zarr_container_path) = {type(zarr_container_path)}."
  )
- file_path_of_element = zarr_container_path / element_type / element_name
- self._validate_can_safely_write_to_path(
- file_path=file_path_of_element, overwrite=overwrite, saving_an_element=True
- )
+ if isinstance(zarr_container_path, Path):
+ file_path_of_element = zarr_container_path / element_type / element_name
+ self._validate_can_safely_write_to_path(
+ file_path=file_path_of_element, overwrite=overwrite, saving_an_element=True
+ )
 
  root_group, element_type_group, _ = self._get_groups_for_element(
  zarr_path=zarr_container_path, element_type=element_type, element_name=element_name
@@ -1376,7 +1381,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.")
@@ -1417,7 +1422,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()
@@ -1438,15 +1443,24 @@ 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
- # 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")
+ from spatialdata._io._utils import _open_zarr_store
+
+ store = _open_zarr_store(self.path)
+ # Note that the store can be local (which does not have the zmetadata bug)
+ # or a remote FSStore (which has the bug).
+ # Consolidate metadata to more easily support remote reading bug in zarr.
+ # We write 'zmetadata' instead of the standard '.zmetadata' to avoid the FSStore bug.
+ # See discussion https://github.com/zarr-developers/zarr-python/issues/1121
+ zarr.consolidate_metadata(store, metadata_key="zmetadata")
  store.close()
 
  def has_consolidated_metadata(self) -> bool:
+ from spatialdata._io._utils import _open_zarr_store
+
  return_value = False
- store = parse_url(self.path, mode="r").store
+ store = _open_zarr_store(self.path)
+ # Note that the store can be local (which does not have the zmetadata bug)
+ # or a remote FSStore (which has the bug).
  if "zmetadata" in store:
  return_value = True
  store.close()
@@ -1575,15 +1589,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

@@ -15,14 +15,18 @@
 from pathlib import Path
 from typing import Any, Literal
 
-import zarr
+import zarr.storage
 from anndata import AnnData
 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._types import StoreLike
 from spatialdata._utils import get_pyramid_levels
 from spatialdata.models._utils import (
  MappingToCoordinateSystem_t,
@@ -31,10 +35,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
@@ -388,6 +389,48 @@ def save_transformations(sdata: SpatialData) -> None:
  sdata.write_transformations()
 
 
+def _open_zarr_store(path: StoreLike, **kwargs: Any) -> zarr.storage.BaseStore:
+ # TODO: ensure kwargs like mode are enforced everywhere and passed correctly to the store
+ if isinstance(path, str | Path):
+ # if the input is str or Path, map it to UPath
+ path = UPath(path)
+ if isinstance(path, PosixUPath | WindowsUPath):
+ # if the input is a local path, use DirectoryStore
+ return zarr.storage.DirectoryStore(path.path, dimension_separator="/")
+ if isinstance(path, zarr.Group):
+ # if the input is a zarr.Group, wrap it with a store
+ if isinstance(path.store, zarr.storage.DirectoryStore):
+ # create a simple FSStore if the store is a DirectoryStore with just the path
+ return FSStore(os.path.join(path.store.path, path.path), **kwargs)
+ if isinstance(path.store, FSStore):
+ # if the store within the zarr.Group is an FSStore, return it
+ # but extend the path of the store with that of the zarr.Group
+ return FSStore(path.store.path + "/" + path.path, fs=path.store.fs, **kwargs)
+ if isinstance(path.store, zarr.storage.ConsolidatedMetadataStore):
+ # if the store is a ConsolidatedMetadataStore, just return the underlying FSSpec store
+ return path.store.store
+ raise ValueError(f"Unsupported store type or zarr.Group: {type(path.store)}")
+ if isinstance(path, zarr.storage.StoreLike):
+ # if the input already a store, wrap it in an FSStore
+ return FSStore(path, **kwargs)
+ if isinstance(path, UPath):
+ # if input is a remote UPath, map it to an FSStore
+ return FSStore(path.path, fs=path.fs, **kwargs)
+ raise TypeError(f"Unsupported type: {type(path)}")
+
+
+def _create_upath(path: StoreLike) -> UPath | None:
+ # try to create a UPath from the input
+ if isinstance(path, zarr.storage.ConsolidatedMetadataStore):
+ path = path.store # get the fsstore from the consolidated store
+ if isinstance(path, FSStore):
+ protocol = path.fs.protocol if isinstance(path.fs.protocol, str) else path.fs.protocol[0]
+ return UPath(path.path, protocol=protocol, **path.fs.storage_options)
+ if isinstance(path, zarr.storage.DirectoryStore):
+ return UPath(path.path)
+ return None
+
+
 class BadFileHandleMethod(Enum):
  ERROR = "error"
  WARN = "warn"

src/spatialdata/_io/format.py

@@ -257,7 +257,7 @@ def validate_table(
 def format_implementations() -> Iterator[Format]:
  """Return an instance of each format implementation, newest to oldest."""
  yield RasterFormatV02()
-# yield RasterFormatV01() # same format string as FormatV04
+ yield RasterFormatV01() # same format string as FormatV04
  yield FormatV04()
  yield FormatV03()
  yield FormatV02()

src/spatialdata/_io/io_points.py

@@ -14,27 +14,19 @@
 )
 from spatialdata._io.format import CurrentPointsFormat, PointsFormats, _parse_version
 from spatialdata.models import get_axes_names
-from spatialdata.transformations._utils import (
- _get_transformations,
- _set_transformations,
-)
+from spatialdata.transformations._utils import _get_transformations, _set_transformations
 
 
 def _read_points(
  store: str | Path | MutableMapping | zarr.Group, # type: ignore[type-arg]
 ) -> DaskDataFrame:
  """Read points from a zarr store."""
- assert isinstance(store, str | Path)
- f = zarr.open(store, mode="r")
-
+ f = zarr.open(store, mode="r") if isinstance(store, str | Path | MutableMapping) else store
  version = _parse_version(f, expect_attrs_key=True)
  assert version is not None
  format = PointsFormats[version]
 
- path = os.path.join(f._store.path, f.path, "points.parquet")
- # cache on remote file needed for parquet reader to work
- # TODO: allow reading in the metadata without caching all the data
- points = read_parquet("simplecache::" + path if path.startswith("http") else path)
+ points = read_parquet(f.store.path, filesystem=getattr(f.store, "fs", None))
  assert isinstance(points, DaskDataFrame)
 
  transformations = _get_transformations_from_ngff_dict(f.attrs.asdict()["coordinateTransformations"])
@@ -57,7 +49,8 @@ def write_points(
  t = _get_transformations(points)
 
  points_groups = group.require_group(name)
- path = Path(points_groups._store.path) / points_groups.path / "points.parquet"
+ store = points_groups._store
+ new_path = os.path.join(store.path, points_groups.path, "points.parquet")
 
  # The following code iterates through all columns in the 'points' DataFrame. If the column's datatype is
  # 'category', it checks whether the categories of this column are known. If not, it explicitly converts the
@@ -70,7 +63,7 @@ def write_points(
  c = c.cat.as_known()
  points[column_name] = c
 
- points.to_parquet(path)
+ points.to_parquet(new_path, filesystem=getattr(store, "fs", None))
 
  attrs = format.attrs_to_dict(points.attrs)
  attrs["version"] = format.spatialdata_format_version

pyproject.toml

@@ -26,20 +26,22 @@ dependencies = [
  "click",
  "dask-image",
  "dask>=2024.4.1,<=2024.11.2",
- "fsspec",
+ "fsspec[s3,http]",
  "geopandas>=0.14",
  "multiscale_spatial_image>=2.0.2",
  "networkx",
  "numba>=0.55.0",
  "numpy",
  "ome_zarr>=0.8.4",
+ "universal_pathlib>=0.2.6",
  "pandas",
  "pooch",
  "pyarrow",
  "rich",
  "setuptools",
  "shapely>=2.0.1",
- "spatial_image>=1.1.0",
+ "spatial_image>=1.2.1",
+ "xarray-dataclasses>=1.9.1",
  "scikit-image",
  "scipy",
  "typing_extensions>=4.8.0",
@@ -58,6 +60,7 @@ test = [
  "pytest-cov",
  "pytest-mock",
  "torch",
+ "moto[s3,server]"
 ]
 docs = [
  "sphinx>=4.5",

src/spatialdata/_core/spatialdata.py

@@ -30,11 +30,8 @@
  validate_table_attr_keys,
 )
 from spatialdata._logging import logger
-from spatialdata._types import ArrayLike, Raster_T
-from spatialdata._utils import (
- _deprecation_alias,
- _error_message_add_element,
-)
+from spatialdata._types import ArrayLike, Raster_T, StoreLike
+from spatialdata._utils import _deprecation_alias, _error_message_add_element
 from spatialdata.models import (
  Image2DModel,
  Image3DModel,
@@ -601,7 +598,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: StoreLike, 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.
@@ -621,9 +618,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}")
@@ -1068,9 +1065,12 @@ def elements_paths_on_disk(self) -> list[str]:
  -------
  A list of paths of the elements saved in the Zarr store.
  """
+ from spatialdata._io._utils import _open_zarr_store
+
  if self.path is None:
  raise ValueError("The SpatialData object is not backed by a Zarr store.")
- store = parse_url(self.path, mode="r").store
+
+ store = _open_zarr_store(self.path)
  root = zarr.group(store=store)
  elements_in_zarr = []
 
@@ -1175,7 +1175,7 @@ def _validate_all_elements(self) -> None:
 
  def write(
  self,
- file_path: str | Path,
+ file_path: StoreLike,
  overwrite: bool = False,
  consolidate_metadata: bool = True,
  format: SpatialDataFormat | list[SpatialDataFormat] | None = None,
@@ -1205,12 +1205,16 @@ def write(
  :class:`~spatialdata._io.format.CurrentRasterFormat`, :class:`~spatialdata._io.format.CurrentShapesFormat`,
  :class:`~spatialdata._io.format.CurrentPointsFormat`, :class:`~spatialdata._io.format.CurrentTablesFormat`.
  """
+ from spatialdata._io._utils import _open_zarr_store
+
  if isinstance(file_path, str):
  file_path = Path(file_path)
- self._validate_can_safely_write_to_path(file_path, overwrite=overwrite)
- self._validate_all_elements()
+ if isinstance(file_path, Path):
+ # TODO: also validate remote paths
+ self._validate_can_safely_write_to_path(file_path, overwrite=overwrite)
+ self._validate_all_elements()
 
- store = parse_url(file_path, mode="w").store
+ store = _open_zarr_store(file_path, mode="w")
  zarr_group = zarr.group(store=store, overwrite=overwrite)
  self.write_attrs(zarr_group=zarr_group)
  store.close()
@@ -1236,20 +1240,21 @@ def write(
  def _write_element(
  self,
  element: SpatialElement | AnnData,
- zarr_container_path: Path,
+ zarr_container_path: StoreLike,
  element_type: str,
  element_name: str,
  overwrite: bool,
  format: SpatialDataFormat | list[SpatialDataFormat] | None = None,
  ) -> None:
- if not isinstance(zarr_container_path, Path):
+ if not isinstance(zarr_container_path, StoreLike):
  raise ValueError(
  f"zarr_container_path must be a Path object, type(zarr_container_path) = {type(zarr_container_path)}."
  )
- file_path_of_element = zarr_container_path / element_type / element_name
- self._validate_can_safely_write_to_path(
- file_path=file_path_of_element, overwrite=overwrite, saving_an_element=True
- )
+ if isinstance(zarr_container_path, Path):
+ file_path_of_element = zarr_container_path / element_type / element_name
+ self._validate_can_safely_write_to_path(
+ file_path=file_path_of_element, overwrite=overwrite, saving_an_element=True
+ )
 
  root_group, element_type_group, _ = self._get_groups_for_element(
  zarr_path=zarr_container_path, element_type=element_type, element_name=element_name
@@ -1376,7 +1381,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.")
@@ -1417,7 +1422,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()
@@ -1438,15 +1443,24 @@ 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
- # 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")
+ from spatialdata._io._utils import _open_zarr_store
+
+ store = _open_zarr_store(self.path)
+ # Note that the store can be local (which does not have the zmetadata bug)
+ # or a remote FSStore (which has the bug).
+ # Consolidate metadata to more easily support remote reading bug in zarr.
+ # We write 'zmetadata' instead of the standard '.zmetadata' to avoid the FSStore bug.
+ # See discussion https://github.com/zarr-developers/zarr-python/issues/1121
+ zarr.consolidate_metadata(store, metadata_key="zmetadata")
  store.close()
 
  def has_consolidated_metadata(self) -> bool:
+ from spatialdata._io._utils import _open_zarr_store
+
  return_value = False
- store = parse_url(self.path, mode="r").store
+ store = _open_zarr_store(self.path)
+ # Note that the store can be local (which does not have the zmetadata bug)
+ # or a remote FSStore (which has the bug).
  if "zmetadata" in store:
  return_value = True
  store.close()
@@ -1575,15 +1589,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

@@ -15,14 +15,18 @@
 from pathlib import Path
 from typing import Any, Literal
 
-import zarr
+import zarr.storage
 from anndata import AnnData
 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._types import StoreLike
 from spatialdata._utils import get_pyramid_levels
 from spatialdata.models._utils import (
  MappingToCoordinateSystem_t,
@@ -31,10 +35,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
@@ -388,6 +389,48 @@ def save_transformations(sdata: SpatialData) -> None:
  sdata.write_transformations()
 
 
+def _open_zarr_store(path: StoreLike, **kwargs: Any) -> zarr.storage.BaseStore:
+ # TODO: ensure kwargs like mode are enforced everywhere and passed correctly to the store
+ if isinstance(path, str | Path):
+ # if the input is str or Path, map it to UPath
+ path = UPath(path)
+ if isinstance(path, PosixUPath | WindowsUPath):
+ # if the input is a local path, use DirectoryStore
+ return zarr.storage.DirectoryStore(path.path, dimension_separator="/")
+ if isinstance(path, zarr.Group):
+ # if the input is a zarr.Group, wrap it with a store
+ if isinstance(path.store, zarr.storage.DirectoryStore):
+ # create a simple FSStore if the store is a DirectoryStore with just the path
+ return FSStore(os.path.join(path.store.path, path.path), **kwargs)
+ if isinstance(path.store, FSStore):
+ # if the store within the zarr.Group is an FSStore, return it
+ # but extend the path of the store with that of the zarr.Group
+ return FSStore(path.store.path + "/" + path.path, fs=path.store.fs, **kwargs)
+ if isinstance(path.store, zarr.storage.ConsolidatedMetadataStore):
+ # if the store is a ConsolidatedMetadataStore, just return the underlying FSSpec store
+ return path.store.store
+ raise ValueError(f"Unsupported store type or zarr.Group: {type(path.store)}")
+ if isinstance(path, zarr.storage.StoreLike):
+ # if the input already a store, wrap it in an FSStore
+ return FSStore(path, **kwargs)
+ if isinstance(path, UPath):
+ # if input is a remote UPath, map it to an FSStore
+ return FSStore(path.path, fs=path.fs, **kwargs)
+ raise TypeError(f"Unsupported type: {type(path)}")
+
+
+def _create_upath(path: StoreLike) -> UPath | None:
+ # try to create a UPath from the input
+ if isinstance(path, zarr.storage.ConsolidatedMetadataStore):
+ path = path.store # get the fsstore from the consolidated store
+ if isinstance(path, FSStore):
+ protocol = path.fs.protocol if isinstance(path.fs.protocol, str) else path.fs.protocol[0]
+ return UPath(path.path, protocol=protocol, **path.fs.storage_options)
+ if isinstance(path, zarr.storage.DirectoryStore):
+ return UPath(path.path)
+ return None
+
+
 class BadFileHandleMethod(Enum):
  ERROR = "error"
  WARN = "warn"

src/spatialdata/_io/format.py

@@ -257,7 +257,7 @@ def validate_table(
 def format_implementations() -> Iterator[Format]:
  """Return an instance of each format implementation, newest to oldest."""
  yield RasterFormatV02()
-# yield RasterFormatV01() # same format string as FormatV04
+ yield RasterFormatV01() # same format string as FormatV04
  yield FormatV04()
  yield FormatV03()
  yield FormatV02()

src/spatialdata/_io/io_points.py

@@ -14,27 +14,19 @@
 )
 from spatialdata._io.format import CurrentPointsFormat, PointsFormats, _parse_version
 from spatialdata.models import get_axes_names
-from spatialdata.transformations._utils import (
- _get_transformations,
- _set_transformations,
-)
+from spatialdata.transformations._utils import _get_transformations, _set_transformations
 
 
 def _read_points(
  store: str | Path | MutableMapping | zarr.Group, # type: ignore[type-arg]
 ) -> DaskDataFrame:
  """Read points from a zarr store."""
- assert isinstance(store, str | Path)
- f = zarr.open(store, mode="r")
-
+ f = zarr.open(store, mode="r") if isinstance(store, str | Path | MutableMapping) else store
  version = _parse_version(f, expect_attrs_key=True)
  assert version is not None
  format = PointsFormats[version]
 
- path = os.path.join(f._store.path, f.path, "points.parquet")
- # cache on remote file needed for parquet reader to work
- # TODO: allow reading in the metadata without caching all the data
- points = read_parquet("simplecache::" + path if path.startswith("http") else path)
+ points = read_parquet(f.store.path, filesystem=getattr(f.store, "fs", None))
  assert isinstance(points, DaskDataFrame)
 
  transformations = _get_transformations_from_ngff_dict(f.attrs.asdict()["coordinateTransformations"])
@@ -57,7 +49,8 @@ def write_points(
  t = _get_transformations(points)
 
  points_groups = group.require_group(name)
- path = Path(points_groups._store.path) / points_groups.path / "points.parquet"
+ store = points_groups._store
+ new_path = os.path.join(store.path, points_groups.path, "points.parquet")
 
  # The following code iterates through all columns in the 'points' DataFrame. If the column's datatype is
  # 'category', it checks whether the categories of this column are known. If not, it explicitly converts the
@@ -70,7 +63,7 @@ def write_points(
  c = c.cat.as_known()
  points[column_name] = c
 
- points.to_parquet(path)
+ points.to_parquet(new_path, filesystem=getattr(store, "fs", None))
 
  attrs = format.attrs_to_dict(points.attrs)
  attrs["version"] = format.spatialdata_format_version

pyproject.toml

@@ -26,20 +26,22 @@ dependencies = [
  "click",
  "dask-image",
  "dask>=2024.4.1,<=2024.11.2",
- "fsspec",
+ "fsspec[s3,http]",
  "geopandas>=0.14",
  "multiscale_spatial_image>=2.0.2",
  "networkx",
  "numba>=0.55.0",
  "numpy",
  "ome_zarr>=0.8.4",
+ "universal_pathlib>=0.2.6",
  "pandas",
  "pooch",
  "pyarrow",
  "rich",
  "setuptools",
  "shapely>=2.0.1",
- "spatial_image>=1.1.0",
+ "spatial_image>=1.2.1",
+ "xarray-dataclasses>=1.9.1",
  "scikit-image",
  "scipy",
  "typing_extensions>=4.8.0",
@@ -58,6 +60,7 @@ test = [
  "pytest-cov",
  "pytest-mock",
  "torch",
+ "moto[s3,server]"
 ]
 docs = [
  "sphinx>=4.5",

src/spatialdata/_core/spatialdata.py

@@ -30,11 +30,8 @@
  validate_table_attr_keys,
 )
 from spatialdata._logging import logger
-from spatialdata._types import ArrayLike, Raster_T
-from spatialdata._utils import (
- _deprecation_alias,
- _error_message_add_element,
-)
+from spatialdata._types import ArrayLike, Raster_T, StoreLike
+from spatialdata._utils import _deprecation_alias, _error_message_add_element
 from spatialdata.models import (
  Image2DModel,
  Image3DModel,
@@ -601,7 +598,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: StoreLike, 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.
@@ -621,9 +618,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}")
@@ -1068,9 +1065,12 @@ def elements_paths_on_disk(self) -> list[str]:
  -------
  A list of paths of the elements saved in the Zarr store.
  """
+ from spatialdata._io._utils import _open_zarr_store
+
  if self.path is None:
  raise ValueError("The SpatialData object is not backed by a Zarr store.")
- store = parse_url(self.path, mode="r").store
+
+ store = _open_zarr_store(self.path)
  root = zarr.group(store=store)
  elements_in_zarr = []
 
@@ -1175,7 +1175,7 @@ def _validate_all_elements(self) -> None:
 
  def write(
  self,
- file_path: str | Path,
+ file_path: StoreLike,
  overwrite: bool = False,
  consolidate_metadata: bool = True,
  format: SpatialDataFormat | list[SpatialDataFormat] | None = None,
@@ -1205,12 +1205,16 @@ def write(
  :class:`~spatialdata._io.format.CurrentRasterFormat`, :class:`~spatialdata._io.format.CurrentShapesFormat`,
  :class:`~spatialdata._io.format.CurrentPointsFormat`, :class:`~spatialdata._io.format.CurrentTablesFormat`.
  """
+ from spatialdata._io._utils import _open_zarr_store
+
  if isinstance(file_path, str):
  file_path = Path(file_path)
- self._validate_can_safely_write_to_path(file_path, overwrite=overwrite)
- self._validate_all_elements()
+ if isinstance(file_path, Path):
+ # TODO: also validate remote paths
+ self._validate_can_safely_write_to_path(file_path, overwrite=overwrite)
+ self._validate_all_elements()
 
- store = parse_url(file_path, mode="w").store
+ store = _open_zarr_store(file_path, mode="w")
  zarr_group = zarr.group(store=store, overwrite=overwrite)
  self.write_attrs(zarr_group=zarr_group)
  store.close()
@@ -1236,20 +1240,21 @@ def write(
  def _write_element(
  self,
  element: SpatialElement | AnnData,
- zarr_container_path: Path,
+ zarr_container_path: StoreLike,
  element_type: str,
  element_name: str,
  overwrite: bool,
  format: SpatialDataFormat | list[SpatialDataFormat] | None = None,
  ) -> None:
- if not isinstance(zarr_container_path, Path):
+ if not isinstance(zarr_container_path, StoreLike):
  raise ValueError(
  f"zarr_container_path must be a Path object, type(zarr_container_path) = {type(zarr_container_path)}."
  )
- file_path_of_element = zarr_container_path / element_type / element_name
- self._validate_can_safely_write_to_path(
- file_path=file_path_of_element, overwrite=overwrite, saving_an_element=True
- )
+ if isinstance(zarr_container_path, Path):
+ file_path_of_element = zarr_container_path / element_type / element_name
+ self._validate_can_safely_write_to_path(
+ file_path=file_path_of_element, overwrite=overwrite, saving_an_element=True
+ )
 
  root_group, element_type_group, _ = self._get_groups_for_element(
  zarr_path=zarr_container_path, element_type=element_type, element_name=element_name
@@ -1376,7 +1381,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.")
@@ -1417,7 +1422,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()
@@ -1438,15 +1443,24 @@ 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
- # 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")
+ from spatialdata._io._utils import _open_zarr_store
+
+ store = _open_zarr_store(self.path)
+ # Note that the store can be local (which does not have the zmetadata bug)
+ # or a remote FSStore (which has the bug).
+ # Consolidate metadata to more easily support remote reading bug in zarr.
+ # We write 'zmetadata' instead of the standard '.zmetadata' to avoid the FSStore bug.
+ # See discussion https://github.com/zarr-developers/zarr-python/issues/1121
+ zarr.consolidate_metadata(store, metadata_key="zmetadata")
  store.close()
 
  def has_consolidated_metadata(self) -> bool:
+ from spatialdata._io._utils import _open_zarr_store
+
  return_value = False
- store = parse_url(self.path, mode="r").store
+ store = _open_zarr_store(self.path)
+ # Note that the store can be local (which does not have the zmetadata bug)
+ # or a remote FSStore (which has the bug).
  if "zmetadata" in store:
  return_value = True
  store.close()
@@ -1575,15 +1589,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

@@ -15,14 +15,18 @@
 from pathlib import Path
 from typing import Any, Literal
 
-import zarr
+import zarr.storage
 from anndata import AnnData
 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._types import StoreLike
 from spatialdata._utils import get_pyramid_levels
 from spatialdata.models._utils import (
  MappingToCoordinateSystem_t,
@@ -31,10 +35,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
@@ -388,6 +389,48 @@ def save_transformations(sdata: SpatialData) -> None:
  sdata.write_transformations()
 
 
+def _open_zarr_store(path: StoreLike, **kwargs: Any) -> zarr.storage.BaseStore:
+ # TODO: ensure kwargs like mode are enforced everywhere and passed correctly to the store
+ if isinstance(path, str | Path):
+ # if the input is str or Path, map it to UPath
+ path = UPath(path)
+ if isinstance(path, PosixUPath | WindowsUPath):
+ # if the input is a local path, use DirectoryStore
+ return zarr.storage.DirectoryStore(path.path, dimension_separator="/")
+ if isinstance(path, zarr.Group):
+ # if the input is a zarr.Group, wrap it with a store
+ if isinstance(path.store, zarr.storage.DirectoryStore):
+ # create a simple FSStore if the store is a DirectoryStore with just the path
+ return FSStore(os.path.join(path.store.path, path.path), **kwargs)
+ if isinstance(path.store, FSStore):
+ # if the store within the zarr.Group is an FSStore, return it
+ # but extend the path of the store with that of the zarr.Group
+ return FSStore(path.store.path + "/" + path.path, fs=path.store.fs, **kwargs)
+ if isinstance(path.store, zarr.storage.ConsolidatedMetadataStore):
+ # if the store is a ConsolidatedMetadataStore, just return the underlying FSSpec store
+ return path.store.store
+ raise ValueError(f"Unsupported store type or zarr.Group: {type(path.store)}")
+ if isinstance(path, zarr.storage.StoreLike):
+ # if the input already a store, wrap it in an FSStore
+ return FSStore(path, **kwargs)
+ if isinstance(path, UPath):
+ # if input is a remote UPath, map it to an FSStore
+ return FSStore(path.path, fs=path.fs, **kwargs)
+ raise TypeError(f"Unsupported type: {type(path)}")
+
+
+def _create_upath(path: StoreLike) -> UPath | None:
+ # try to create a UPath from the input
+ if isinstance(path, zarr.storage.ConsolidatedMetadataStore):
+ path = path.store # get the fsstore from the consolidated store
+ if isinstance(path, FSStore):
+ protocol = path.fs.protocol if isinstance(path.fs.protocol, str) else path.fs.protocol[0]
+ return UPath(path.path, protocol=protocol, **path.fs.storage_options)
+ if isinstance(path, zarr.storage.DirectoryStore):
+ return UPath(path.path)
+ return None
+
+
 class BadFileHandleMethod(Enum):
  ERROR = "error"
  WARN = "warn"

src/spatialdata/_io/format.py

@@ -257,7 +257,7 @@ def validate_table(
 def format_implementations() -> Iterator[Format]:
  """Return an instance of each format implementation, newest to oldest."""
  yield RasterFormatV02()
-# yield RasterFormatV01() # same format string as FormatV04
+ yield RasterFormatV01() # same format string as FormatV04
  yield FormatV04()
  yield FormatV03()
  yield FormatV02()

src/spatialdata/_io/io_points.py

@@ -14,27 +14,19 @@
 )
 from spatialdata._io.format import CurrentPointsFormat, PointsFormats, _parse_version
 from spatialdata.models import get_axes_names
-from spatialdata.transformations._utils import (
- _get_transformations,
- _set_transformations,
-)
+from spatialdata.transformations._utils import _get_transformations, _set_transformations
 
 
 def _read_points(
  store: str | Path | MutableMapping | zarr.Group, # type: ignore[type-arg]
 ) -> DaskDataFrame:
  """Read points from a zarr store."""
- assert isinstance(store, str | Path)
- f = zarr.open(store, mode="r")
-
+ f = zarr.open(store, mode="r") if isinstance(store, str | Path | MutableMapping) else store
  version = _parse_version(f, expect_attrs_key=True)
  assert version is not None
  format = PointsFormats[version]
 
- path = os.path.join(f._store.path, f.path, "points.parquet")
- # cache on remote file needed for parquet reader to work
- # TODO: allow reading in the metadata without caching all the data
- points = read_parquet("simplecache::" + path if path.startswith("http") else path)
+ points = read_parquet(f.store.path, filesystem=getattr(f.store, "fs", None))
  assert isinstance(points, DaskDataFrame)
 
  transformations = _get_transformations_from_ngff_dict(f.attrs.asdict()["coordinateTransformations"])
@@ -57,7 +49,8 @@ def write_points(
  t = _get_transformations(points)
 
  points_groups = group.require_group(name)
- path = Path(points_groups._store.path) / points_groups.path / "points.parquet"
+ store = points_groups._store
+ new_path = os.path.join(store.path, points_groups.path, "points.parquet")
 
  # The following code iterates through all columns in the 'points' DataFrame. If the column's datatype is
  # 'category', it checks whether the categories of this column are known. If not, it explicitly converts the
@@ -70,7 +63,7 @@ def write_points(
  c = c.cat.as_known()
  points[column_name] = c
 
- points.to_parquet(path)
+ points.to_parquet(new_path, filesystem=getattr(store, "fs", None))
 
  attrs = format.attrs_to_dict(points.attrs)
  attrs["version"] = format.spatialdata_format_version

pyproject.toml

@@ -26,20 +26,22 @@ dependencies = [
  "click",
  "dask-image",
  "dask>=2024.4.1,<=2024.11.2",
- "fsspec",
+ "fsspec[s3,http]",
  "geopandas>=0.14",
  "multiscale_spatial_image>=2.0.2",
  "networkx",
  "numba>=0.55.0",
  "numpy",
  "ome_zarr>=0.8.4",
+ "universal_pathlib>=0.2.6",
  "pandas",
  "pooch",
  "pyarrow",
  "rich",
  "setuptools",
  "shapely>=2.0.1",
- "spatial_image>=1.1.0",
+ "spatial_image>=1.2.1",
+ "xarray-dataclasses>=1.9.1",
  "scikit-image",
  "scipy",
  "typing_extensions>=4.8.0",
@@ -58,6 +60,7 @@ test = [
  "pytest-cov",
  "pytest-mock",
  "torch",
+ "moto[s3,server]"
 ]
 docs = [
  "sphinx>=4.5",

src/spatialdata/_core/spatialdata.py

@@ -30,11 +30,8 @@
  validate_table_attr_keys,
 )
 from spatialdata._logging import logger
-from spatialdata._types import ArrayLike, Raster_T
-from spatialdata._utils import (
- _deprecation_alias,
- _error_message_add_element,
-)
+from spatialdata._types import ArrayLike, Raster_T, StoreLike
+from spatialdata._utils import _deprecation_alias, _error_message_add_element
 from spatialdata.models import (
  Image2DModel,
  Image3DModel,
@@ -601,7 +598,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: StoreLike, 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.
@@ -621,9 +618,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}")
@@ -1068,9 +1065,12 @@ def elements_paths_on_disk(self) -> list[str]:
  -------
  A list of paths of the elements saved in the Zarr store.
  """
+ from spatialdata._io._utils import _open_zarr_store
+
  if self.path is None:
  raise ValueError("The SpatialData object is not backed by a Zarr store.")
- store = parse_url(self.path, mode="r").store
+
+ store = _open_zarr_store(self.path)
  root = zarr.group(store=store)
  elements_in_zarr = []
 
@@ -1175,7 +1175,7 @@ def _validate_all_elements(self) -> None:
 
  def write(
  self,
- file_path: str | Path,
+ file_path: StoreLike,
  overwrite: bool = False,
  consolidate_metadata: bool = True,
  format: SpatialDataFormat | list[SpatialDataFormat] | None = None,
@@ -1205,12 +1205,16 @@ def write(
  :class:`~spatialdata._io.format.CurrentRasterFormat`, :class:`~spatialdata._io.format.CurrentShapesFormat`,
  :class:`~spatialdata._io.format.CurrentPointsFormat`, :class:`~spatialdata._io.format.CurrentTablesFormat`.
  """
+ from spatialdata._io._utils import _open_zarr_store
+
  if isinstance(file_path, str):
  file_path = Path(file_path)
- self._validate_can_safely_write_to_path(file_path, overwrite=overwrite)
- self._validate_all_elements()
+ if isinstance(file_path, Path):
+ # TODO: also validate remote paths
+ self._validate_can_safely_write_to_path(file_path, overwrite=overwrite)
+ self._validate_all_elements()
 
- store = parse_url(file_path, mode="w").store
+ store = _open_zarr_store(file_path, mode="w")
  zarr_group = zarr.group(store=store, overwrite=overwrite)
  self.write_attrs(zarr_group=zarr_group)
  store.close()
@@ -1236,20 +1240,21 @@ def write(
  def _write_element(
  self,
  element: SpatialElement | AnnData,
- zarr_container_path: Path,
+ zarr_container_path: StoreLike,
  element_type: str,
  element_name: str,
  overwrite: bool,
  format: SpatialDataFormat | list[SpatialDataFormat] | None = None,
  ) -> None:
- if not isinstance(zarr_container_path, Path):
+ if not isinstance(zarr_container_path, StoreLike):
  raise ValueError(
  f"zarr_container_path must be a Path object, type(zarr_container_path) = {type(zarr_container_path)}."
  )
- file_path_of_element = zarr_container_path / element_type / element_name
- self._validate_can_safely_write_to_path(
- file_path=file_path_of_element, overwrite=overwrite, saving_an_element=True
- )
+ if isinstance(zarr_container_path, Path):
+ file_path_of_element = zarr_container_path / element_type / element_name
+ self._validate_can_safely_write_to_path(
+ file_path=file_path_of_element, overwrite=overwrite, saving_an_element=True
+ )
 
  root_group, element_type_group, _ = self._get_groups_for_element(
  zarr_path=zarr_container_path, element_type=element_type, element_name=element_name
@@ -1376,7 +1381,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.")
@@ -1417,7 +1422,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()
@@ -1438,15 +1443,24 @@ 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
- # 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")
+ from spatialdata._io._utils import _open_zarr_store
+
+ store = _open_zarr_store(self.path)
+ # Note that the store can be local (which does not have the zmetadata bug)
+ # or a remote FSStore (which has the bug).
+ # Consolidate metadata to more easily support remote reading bug in zarr.
+ # We write 'zmetadata' instead of the standard '.zmetadata' to avoid the FSStore bug.
+ # See discussion https://github.com/zarr-developers/zarr-python/issues/1121
+ zarr.consolidate_metadata(store, metadata_key="zmetadata")
  store.close()
 
  def has_consolidated_metadata(self) -> bool:
+ from spatialdata._io._utils import _open_zarr_store
+
  return_value = False
- store = parse_url(self.path, mode="r").store
+ store = _open_zarr_store(self.path)
+ # Note that the store can be local (which does not have the zmetadata bug)
+ # or a remote FSStore (which has the bug).
  if "zmetadata" in store:
  return_value = True
  store.close()
@@ -1575,15 +1589,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

@@ -15,14 +15,18 @@
 from pathlib import Path
 from typing import Any, Literal
 
-import zarr
+import zarr.storage
 from anndata import AnnData
 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._types import StoreLike
 from spatialdata._utils import get_pyramid_levels
 from spatialdata.models._utils import (
  MappingToCoordinateSystem_t,
@@ -31,10 +35,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
@@ -388,6 +389,48 @@ def save_transformations(sdata: SpatialData) -> None:
  sdata.write_transformations()
 
 
+def _open_zarr_store(path: StoreLike, **kwargs: Any) -> zarr.storage.BaseStore:
+ # TODO: ensure kwargs like mode are enforced everywhere and passed correctly to the store
+ if isinstance(path, str | Path):
+ # if the input is str or Path, map it to UPath
+ path = UPath(path)
+ if isinstance(path, PosixUPath | WindowsUPath):
+ # if the input is a local path, use DirectoryStore
+ return zarr.storage.DirectoryStore(path.path, dimension_separator="/")
+ if isinstance(path, zarr.Group):
+ # if the input is a zarr.Group, wrap it with a store
+ if isinstance(path.store, zarr.storage.DirectoryStore):
+ # create a simple FSStore if the store is a DirectoryStore with just the path
+ return FSStore(os.path.join(path.store.path, path.path), **kwargs)
+ if isinstance(path.store, FSStore):
+ # if the store within the zarr.Group is an FSStore, return it
+ # but extend the path of the store with that of the zarr.Group
+ return FSStore(path.store.path + "/" + path.path, fs=path.store.fs, **kwargs)
+ if isinstance(path.store, zarr.storage.ConsolidatedMetadataStore):
+ # if the store is a ConsolidatedMetadataStore, just return the underlying FSSpec store
+ return path.store.store
+ raise ValueError(f"Unsupported store type or zarr.Group: {type(path.store)}")
+ if isinstance(path, zarr.storage.StoreLike):
+ # if the input already a store, wrap it in an FSStore
+ return FSStore(path, **kwargs)
+ if isinstance(path, UPath):
+ # if input is a remote UPath, map it to an FSStore
+ return FSStore(path.path, fs=path.fs, **kwargs)
+ raise TypeError(f"Unsupported type: {type(path)}")
+
+
+def _create_upath(path: StoreLike) -> UPath | None:
+ # try to create a UPath from the input
+ if isinstance(path, zarr.storage.ConsolidatedMetadataStore):
+ path = path.store # get the fsstore from the consolidated store
+ if isinstance(path, FSStore):
+ protocol = path.fs.protocol if isinstance(path.fs.protocol, str) else path.fs.protocol[0]
+ return UPath(path.path, protocol=protocol, **path.fs.storage_options)
+ if isinstance(path, zarr.storage.DirectoryStore):
+ return UPath(path.path)
+ return None
+
+
 class BadFileHandleMethod(Enum):
  ERROR = "error"
  WARN = "warn"

src/spatialdata/_io/format.py

@@ -257,7 +257,7 @@ def validate_table(
 def format_implementations() -> Iterator[Format]:
  """Return an instance of each format implementation, newest to oldest."""
  yield RasterFormatV02()
-# yield RasterFormatV01() # same format string as FormatV04
+ yield RasterFormatV01() # same format string as FormatV04
  yield FormatV04()
  yield FormatV03()
  yield FormatV02()

src/spatialdata/_io/io_points.py

@@ -14,27 +14,19 @@
 )
 from spatialdata._io.format import CurrentPointsFormat, PointsFormats, _parse_version
 from spatialdata.models import get_axes_names
-from spatialdata.transformations._utils import (
- _get_transformations,
- _set_transformations,
-)
+from spatialdata.transformations._utils import _get_transformations, _set_transformations
 
 
 def _read_points(
  store: str | Path | MutableMapping | zarr.Group, # type: ignore[type-arg]
 ) -> DaskDataFrame:
  """Read points from a zarr store."""
- assert isinstance(store, str | Path)
- f = zarr.open(store, mode="r")
-
+ f = zarr.open(store, mode="r") if isinstance(store, str | Path | MutableMapping) else store
  version = _parse_version(f, expect_attrs_key=True)
  assert version is not None
  format = PointsFormats[version]
 
- path = os.path.join(f._store.path, f.path, "points.parquet")
- # cache on remote file needed for parquet reader to work
- # TODO: allow reading in the metadata without caching all the data
- points = read_parquet("simplecache::" + path if path.startswith("http") else path)
+ points = read_parquet(f.store.path, filesystem=getattr(f.store, "fs", None))
  assert isinstance(points, DaskDataFrame)
 
  transformations = _get_transformations_from_ngff_dict(f.attrs.asdict()["coordinateTransformations"])
@@ -57,7 +49,8 @@ def write_points(
  t = _get_transformations(points)
 
  points_groups = group.require_group(name)
- path = Path(points_groups._store.path) / points_groups.path / "points.parquet"
+ store = points_groups._store
+ new_path = os.path.join(store.path, points_groups.path, "points.parquet")
 
  # The following code iterates through all columns in the 'points' DataFrame. If the column's datatype is
  # 'category', it checks whether the categories of this column are known. If not, it explicitly converts the
@@ -70,7 +63,7 @@ def write_points(
  c = c.cat.as_known()
  points[column_name] = c
 
- points.to_parquet(path)
+ points.to_parquet(new_path, filesystem=getattr(store, "fs", None))
 
  attrs = format.attrs_to_dict(points.attrs)
  attrs["version"] = format.spatialdata_format_version

pyproject.toml

@@ -26,20 +26,22 @@ dependencies = [
  "click",
  "dask-image",
  "dask>=2024.4.1,<=2024.11.2",
- "fsspec",
+ "fsspec[s3,http]",
  "geopandas>=0.14",
  "multiscale_spatial_image>=2.0.2",
  "networkx",
  "numba>=0.55.0",
  "numpy",
  "ome_zarr>=0.8.4",
+ "universal_pathlib>=0.2.6",
  "pandas",
  "pooch",
  "pyarrow",
  "rich",
  "setuptools",
  "shapely>=2.0.1",
- "spatial_image>=1.1.0",
+ "spatial_image>=1.2.1",
+ "xarray-dataclasses>=1.9.1",
  "scikit-image",
  "scipy",
  "typing_extensions>=4.8.0",
@@ -58,6 +60,7 @@ test = [
  "pytest-cov",
  "pytest-mock",
  "torch",
+ "moto[s3,server]"
 ]
 docs = [
  "sphinx>=4.5",

src/spatialdata/_core/spatialdata.py

@@ -30,11 +30,8 @@
  validate_table_attr_keys,
 )
 from spatialdata._logging import logger
-from spatialdata._types import ArrayLike, Raster_T
-from spatialdata._utils import (
- _deprecation_alias,
- _error_message_add_element,
-)
+from spatialdata._types import ArrayLike, Raster_T, StoreLike
+from spatialdata._utils import _deprecation_alias, _error_message_add_element
 from spatialdata.models import (
  Image2DModel,
  Image3DModel,
@@ -601,7 +598,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: StoreLike, 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.
@@ -621,9 +618,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}")
@@ -1068,9 +1065,12 @@ def elements_paths_on_disk(self) -> list[str]:
  -------
  A list of paths of the elements saved in the Zarr store.
  """
+ from spatialdata._io._utils import _open_zarr_store
+
  if self.path is None:
  raise ValueError("The SpatialData object is not backed by a Zarr store.")
- store = parse_url(self.path, mode="r").store
+
+ store = _open_zarr_store(self.path)
  root = zarr.group(store=store)
  elements_in_zarr = []
 
@@ -1175,7 +1175,7 @@ def _validate_all_elements(self) -> None:
 
  def write(
  self,
- file_path: str | Path,
+ file_path: StoreLike,
  overwrite: bool = False,
  consolidate_metadata: bool = True,
  format: SpatialDataFormat | list[SpatialDataFormat] | None = None,
@@ -1205,12 +1205,16 @@ def write(
  :class:`~spatialdata._io.format.CurrentRasterFormat`, :class:`~spatialdata._io.format.CurrentShapesFormat`,
  :class:`~spatialdata._io.format.CurrentPointsFormat`, :class:`~spatialdata._io.format.CurrentTablesFormat`.
  """
+ from spatialdata._io._utils import _open_zarr_store
+
  if isinstance(file_path, str):
  file_path = Path(file_path)
- self._validate_can_safely_write_to_path(file_path, overwrite=overwrite)
- self._validate_all_elements()
+ if isinstance(file_path, Path):
+ # TODO: also validate remote paths
+ self._validate_can_safely_write_to_path(file_path, overwrite=overwrite)
+ self._validate_all_elements()
 
- store = parse_url(file_path, mode="w").store
+ store = _open_zarr_store(file_path, mode="w")
  zarr_group = zarr.group(store=store, overwrite=overwrite)
  self.write_attrs(zarr_group=zarr_group)
  store.close()
@@ -1236,20 +1240,21 @@ def write(
  def _write_element(
  self,
  element: SpatialElement | AnnData,
- zarr_container_path: Path,
+ zarr_container_path: StoreLike,
  element_type: str,
  element_name: str,
  overwrite: bool,
  format: SpatialDataFormat | list[SpatialDataFormat] | None = None,
  ) -> None:
- if not isinstance(zarr_container_path, Path):
+ if not isinstance(zarr_container_path, StoreLike):
  raise ValueError(
  f"zarr_container_path must be a Path object, type(zarr_container_path) = {type(zarr_container_path)}."
  )
- file_path_of_element = zarr_container_path / element_type / element_name
- self._validate_can_safely_write_to_path(
- file_path=file_path_of_element, overwrite=overwrite, saving_an_element=True
- )
+ if isinstance(zarr_container_path, Path):
+ file_path_of_element = zarr_container_path / element_type / element_name
+ self._validate_can_safely_write_to_path(
+ file_path=file_path_of_element, overwrite=overwrite, saving_an_element=True
+ )
 
  root_group, element_type_group, _ = self._get_groups_for_element(
  zarr_path=zarr_container_path, element_type=element_type, element_name=element_name
@@ -1376,7 +1381,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.")
@@ -1417,7 +1422,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()
@@ -1438,15 +1443,24 @@ 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
- # 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")
+ from spatialdata._io._utils import _open_zarr_store
+
+ store = _open_zarr_store(self.path)
+ # Note that the store can be local (which does not have the zmetadata bug)
+ # or a remote FSStore (which has the bug).
+ # Consolidate metadata to more easily support remote reading bug in zarr.
+ # We write 'zmetadata' instead of the standard '.zmetadata' to avoid the FSStore bug.
+ # See discussion https://github.com/zarr-developers/zarr-python/issues/1121
+ zarr.consolidate_metadata(store, metadata_key="zmetadata")
  store.close()
 
  def has_consolidated_metadata(self) -> bool:
+ from spatialdata._io._utils import _open_zarr_store
+
  return_value = False
- store = parse_url(self.path, mode="r").store
+ store = _open_zarr_store(self.path)
+ # Note that the store can be local (which does not have the zmetadata bug)
+ # or a remote FSStore (which has the bug).
  if "zmetadata" in store:
  return_value = True
  store.close()
@@ -1575,15 +1589,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

@@ -15,14 +15,18 @@
 from pathlib import Path
 from typing import Any, Literal
 
-import zarr
+import zarr.storage
 from anndata import AnnData
 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._types import StoreLike
 from spatialdata._utils import get_pyramid_levels
 from spatialdata.models._utils import (
  MappingToCoordinateSystem_t,
@@ -31,10 +35,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
@@ -388,6 +389,48 @@ def save_transformations(sdata: SpatialData) -> None:
  sdata.write_transformations()
 
 
+def _open_zarr_store(path: StoreLike, **kwargs: Any) -> zarr.storage.BaseStore:
+ # TODO: ensure kwargs like mode are enforced everywhere and passed correctly to the store
+ if isinstance(path, str | Path):
+ # if the input is str or Path, map it to UPath
+ path = UPath(path)
+ if isinstance(path, PosixUPath | WindowsUPath):
+ # if the input is a local path, use DirectoryStore
+ return zarr.storage.DirectoryStore(path.path, dimension_separator="/")
+ if isinstance(path, zarr.Group):
+ # if the input is a zarr.Group, wrap it with a store
+ if isinstance(path.store, zarr.storage.DirectoryStore):
+ # create a simple FSStore if the store is a DirectoryStore with just the path
+ return FSStore(os.path.join(path.store.path, path.path), **kwargs)
+ if isinstance(path.store, FSStore):
+ # if the store within the zarr.Group is an FSStore, return it
+ # but extend the path of the store with that of the zarr.Group
+ return FSStore(path.store.path + "/" + path.path, fs=path.store.fs, **kwargs)
+ if isinstance(path.store, zarr.storage.ConsolidatedMetadataStore):
+ # if the store is a ConsolidatedMetadataStore, just return the underlying FSSpec store
+ return path.store.store
+ raise ValueError(f"Unsupported store type or zarr.Group: {type(path.store)}")
+ if isinstance(path, zarr.storage.StoreLike):
+ # if the input already a store, wrap it in an FSStore
+ return FSStore(path, **kwargs)
+ if isinstance(path, UPath):
+ # if input is a remote UPath, map it to an FSStore
+ return FSStore(path.path, fs=path.fs, **kwargs)
+ raise TypeError(f"Unsupported type: {type(path)}")
+
+
+def _create_upath(path: StoreLike) -> UPath | None:
+ # try to create a UPath from the input
+ if isinstance(path, zarr.storage.ConsolidatedMetadataStore):
+ path = path.store # get the fsstore from the consolidated store
+ if isinstance(path, FSStore):
+ protocol = path.fs.protocol if isinstance(path.fs.protocol, str) else path.fs.protocol[0]
+ return UPath(path.path, protocol=protocol, **path.fs.storage_options)
+ if isinstance(path, zarr.storage.DirectoryStore):
+ return UPath(path.path)
+ return None
+
+
 class BadFileHandleMethod(Enum):
  ERROR = "error"
  WARN = "warn"

src/spatialdata/_io/format.py

@@ -257,7 +257,7 @@ def validate_table(
 def format_implementations() -> Iterator[Format]:
  """Return an instance of each format implementation, newest to oldest."""
  yield RasterFormatV02()
-# yield RasterFormatV01() # same format string as FormatV04
+ yield RasterFormatV01() # same format string as FormatV04
  yield FormatV04()
  yield FormatV03()
  yield FormatV02()

src/spatialdata/_io/io_points.py

@@ -14,27 +14,19 @@
 )
 from spatialdata._io.format import CurrentPointsFormat, PointsFormats, _parse_version
 from spatialdata.models import get_axes_names
-from spatialdata.transformations._utils import (
- _get_transformations,
- _set_transformations,
-)
+from spatialdata.transformations._utils import _get_transformations, _set_transformations
 
 
 def _read_points(
  store: str | Path | MutableMapping | zarr.Group, # type: ignore[type-arg]
 ) -> DaskDataFrame:
  """Read points from a zarr store."""
- assert isinstance(store, str | Path)
- f = zarr.open(store, mode="r")
-
+ f = zarr.open(store, mode="r") if isinstance(store, str | Path | MutableMapping) else store
  version = _parse_version(f, expect_attrs_key=True)
  assert version is not None
  format = PointsFormats[version]
 
- path = os.path.join(f._store.path, f.path, "points.parquet")
- # cache on remote file needed for parquet reader to work
- # TODO: allow reading in the metadata without caching all the data
- points = read_parquet("simplecache::" + path if path.startswith("http") else path)
+ points = read_parquet(f.store.path, filesystem=getattr(f.store, "fs", None))
  assert isinstance(points, DaskDataFrame)
 
  transformations = _get_transformations_from_ngff_dict(f.attrs.asdict()["coordinateTransformations"])
@@ -57,7 +49,8 @@ def write_points(
  t = _get_transformations(points)
 
  points_groups = group.require_group(name)
- path = Path(points_groups._store.path) / points_groups.path / "points.parquet"
+ store = points_groups._store
+ new_path = os.path.join(store.path, points_groups.path, "points.parquet")
 
  # The following code iterates through all columns in the 'points' DataFrame. If the column's datatype is
  # 'category', it checks whether the categories of this column are known. If not, it explicitly converts the
@@ -70,7 +63,7 @@ def write_points(
  c = c.cat.as_known()
  points[column_name] = c
 
- points.to_parquet(path)
+ points.to_parquet(new_path, filesystem=getattr(store, "fs", None))
 
  attrs = format.attrs_to_dict(points.attrs)
  attrs["version"] = format.spatialdata_format_version

pyproject.toml

@@ -26,20 +26,22 @@ dependencies = [
  "click",
  "dask-image",
  "dask>=2024.4.1,<=2024.11.2",
- "fsspec",
+ "fsspec[s3,http]",
  "geopandas>=0.14",
  "multiscale_spatial_image>=2.0.2",
  "networkx",
  "numba>=0.55.0",
  "numpy",
  "ome_zarr>=0.8.4",
+ "universal_pathlib>=0.2.6",
  "pandas",
  "pooch",
  "pyarrow",
  "rich",
  "setuptools",
  "shapely>=2.0.1",
- "spatial_image>=1.1.0",
+ "spatial_image>=1.2.1",
+ "xarray-dataclasses>=1.9.1",
  "scikit-image",
  "scipy",
  "typing_extensions>=4.8.0",
@@ -58,6 +60,7 @@ test = [
  "pytest-cov",
  "pytest-mock",
  "torch",
+ "moto[s3,server]"
 ]
 docs = [
  "sphinx>=4.5",

src/spatialdata/_core/spatialdata.py

@@ -30,11 +30,8 @@
  validate_table_attr_keys,
 )
 from spatialdata._logging import logger
-from spatialdata._types import ArrayLike, Raster_T
-from spatialdata._utils import (
- _deprecation_alias,
- _error_message_add_element,
-)
+from spatialdata._types import ArrayLike, Raster_T, StoreLike
+from spatialdata._utils import _deprecation_alias, _error_message_add_element
 from spatialdata.models import (
  Image2DModel,
  Image3DModel,
@@ -601,7 +598,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: StoreLike, 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.
@@ -621,9 +618,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}")
@@ -1068,9 +1065,12 @@ def elements_paths_on_disk(self) -> list[str]:
  -------
  A list of paths of the elements saved in the Zarr store.
  """
+ from spatialdata._io._utils import _open_zarr_store
+
  if self.path is None:
  raise ValueError("The SpatialData object is not backed by a Zarr store.")
- store = parse_url(self.path, mode="r").store
+
+ store = _open_zarr_store(self.path)
  root = zarr.group(store=store)
  elements_in_zarr = []
 
@@ -1175,7 +1175,7 @@ def _validate_all_elements(self) -> None:
 
  def write(
  self,
- file_path: str | Path,
+ file_path: StoreLike,
  overwrite: bool = False,
  consolidate_metadata: bool = True,
  format: SpatialDataFormat | list[SpatialDataFormat] | None = None,
@@ -1205,12 +1205,16 @@ def write(
  :class:`~spatialdata._io.format.CurrentRasterFormat`, :class:`~spatialdata._io.format.CurrentShapesFormat`,
  :class:`~spatialdata._io.format.CurrentPointsFormat`, :class:`~spatialdata._io.format.CurrentTablesFormat`.
  """
+ from spatialdata._io._utils import _open_zarr_store
+
  if isinstance(file_path, str):
  file_path = Path(file_path)
- self._validate_can_safely_write_to_path(file_path, overwrite=overwrite)
- self._validate_all_elements()
+ if isinstance(file_path, Path):
+ # TODO: also validate remote paths
+ self._validate_can_safely_write_to_path(file_path, overwrite=overwrite)
+ self._validate_all_elements()
 
- store = parse_url(file_path, mode="w").store
+ store = _open_zarr_store(file_path, mode="w")
  zarr_group = zarr.group(store=store, overwrite=overwrite)
  self.write_attrs(zarr_group=zarr_group)
  store.close()
@@ -1236,20 +1240,21 @@ def write(
  def _write_element(
  self,
  element: SpatialElement | AnnData,
- zarr_container_path: Path,
+ zarr_container_path: StoreLike,
  element_type: str,
  element_name: str,
  overwrite: bool,
  format: SpatialDataFormat | list[SpatialDataFormat] | None = None,
  ) -> None:
- if not isinstance(zarr_container_path, Path):
+ if not isinstance(zarr_container_path, StoreLike):
  raise ValueError(
  f"zarr_container_path must be a Path object, type(zarr_container_path) = {type(zarr_container_path)}."
  )
- file_path_of_element = zarr_container_path / element_type / element_name
- self._validate_can_safely_write_to_path(
- file_path=file_path_of_element, overwrite=overwrite, saving_an_element=True
- )
+ if isinstance(zarr_container_path, Path):
+ file_path_of_element = zarr_container_path / element_type / element_name
+ self._validate_can_safely_write_to_path(
+ file_path=file_path_of_element, overwrite=overwrite, saving_an_element=True
+ )
 
  root_group, element_type_group, _ = self._get_groups_for_element(
  zarr_path=zarr_container_path, element_type=element_type, element_name=element_name
@@ -1376,7 +1381,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.")
@@ -1417,7 +1422,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()
@@ -1438,15 +1443,24 @@ 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
- # 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")
+ from spatialdata._io._utils import _open_zarr_store
+
+ store = _open_zarr_store(self.path)
+ # Note that the store can be local (which does not have the zmetadata bug)
+ # or a remote FSStore (which has the bug).
+ # Consolidate metadata to more easily support remote reading bug in zarr.
+ # We write 'zmetadata' instead of the standard '.zmetadata' to avoid the FSStore bug.
+ # See discussion https://github.com/zarr-developers/zarr-python/issues/1121
+ zarr.consolidate_metadata(store, metadata_key="zmetadata")
  store.close()
 
  def has_consolidated_metadata(self) -> bool:
+ from spatialdata._io._utils import _open_zarr_store
+
  return_value = False
- store = parse_url(self.path, mode="r").store
+ store = _open_zarr_store(self.path)
+ # Note that the store can be local (which does not have the zmetadata bug)
+ # or a remote FSStore (which has the bug).
  if "zmetadata" in store:
  return_value = True
  store.close()
@@ -1575,15 +1589,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

@@ -15,14 +15,18 @@
 from pathlib import Path
 from typing import Any, Literal
 
-import zarr
+import zarr.storage
 from anndata import AnnData
 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._types import StoreLike
 from spatialdata._utils import get_pyramid_levels
 from spatialdata.models._utils import (
  MappingToCoordinateSystem_t,
@@ -31,10 +35,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
@@ -388,6 +389,48 @@ def save_transformations(sdata: SpatialData) -> None:
  sdata.write_transformations()
 
 
+def _open_zarr_store(path: StoreLike, **kwargs: Any) -> zarr.storage.BaseStore:
+ # TODO: ensure kwargs like mode are enforced everywhere and passed correctly to the store
+ if isinstance(path, str | Path):
+ # if the input is str or Path, map it to UPath
+ path = UPath(path)
+ if isinstance(path, PosixUPath | WindowsUPath):
+ # if the input is a local path, use DirectoryStore
+ return zarr.storage.DirectoryStore(path.path, dimension_separator="/")
+ if isinstance(path, zarr.Group):
+ # if the input is a zarr.Group, wrap it with a store
+ if isinstance(path.store, zarr.storage.DirectoryStore):
+ # create a simple FSStore if the store is a DirectoryStore with just the path
+ return FSStore(os.path.join(path.store.path, path.path), **kwargs)
+ if isinstance(path.store, FSStore):
+ # if the store within the zarr.Group is an FSStore, return it
+ # but extend the path of the store with that of the zarr.Group
+ return FSStore(path.store.path + "/" + path.path, fs=path.store.fs, **kwargs)
+ if isinstance(path.store, zarr.storage.ConsolidatedMetadataStore):
+ # if the store is a ConsolidatedMetadataStore, just return the underlying FSSpec store
+ return path.store.store
+ raise ValueError(f"Unsupported store type or zarr.Group: {type(path.store)}")
+ if isinstance(path, zarr.storage.StoreLike):
+ # if the input already a store, wrap it in an FSStore
+ return FSStore(path, **kwargs)
+ if isinstance(path, UPath):
+ # if input is a remote UPath, map it to an FSStore
+ return FSStore(path.path, fs=path.fs, **kwargs)
+ raise TypeError(f"Unsupported type: {type(path)}")
+
+
+def _create_upath(path: StoreLike) -> UPath | None:
+ # try to create a UPath from the input
+ if isinstance(path, zarr.storage.ConsolidatedMetadataStore):
+ path = path.store # get the fsstore from the consolidated store
+ if isinstance(path, FSStore):
+ protocol = path.fs.protocol if isinstance(path.fs.protocol, str) else path.fs.protocol[0]
+ return UPath(path.path, protocol=protocol, **path.fs.storage_options)
+ if isinstance(path, zarr.storage.DirectoryStore):
+ return UPath(path.path)
+ return None
+
+
 class BadFileHandleMethod(Enum):
  ERROR = "error"
  WARN = "warn"

src/spatialdata/_io/format.py

@@ -257,7 +257,7 @@ def validate_table(
 def format_implementations() -> Iterator[Format]:
  """Return an instance of each format implementation, newest to oldest."""
  yield RasterFormatV02()
-# yield RasterFormatV01() # same format string as FormatV04
+ yield RasterFormatV01() # same format string as FormatV04
  yield FormatV04()
  yield FormatV03()
  yield FormatV02()

src/spatialdata/_io/io_points.py

@@ -14,27 +14,19 @@
 )
 from spatialdata._io.format import CurrentPointsFormat, PointsFormats, _parse_version
 from spatialdata.models import get_axes_names
-from spatialdata.transformations._utils import (
- _get_transformations,
- _set_transformations,
-)
+from spatialdata.transformations._utils import _get_transformations, _set_transformations
 
 
 def _read_points(
  store: str | Path | MutableMapping | zarr.Group, # type: ignore[type-arg]
 ) -> DaskDataFrame:
  """Read points from a zarr store."""
- assert isinstance(store, str | Path)
- f = zarr.open(store, mode="r")
-
+ f = zarr.open(store, mode="r") if isinstance(store, str | Path | MutableMapping) else store
  version = _parse_version(f, expect_attrs_key=True)
  assert version is not None
  format = PointsFormats[version]
 
- path = os.path.join(f._store.path, f.path, "points.parquet")
- # cache on remote file needed for parquet reader to work
- # TODO: allow reading in the metadata without caching all the data
- points = read_parquet("simplecache::" + path if path.startswith("http") else path)
+ points = read_parquet(f.store.path, filesystem=getattr(f.store, "fs", None))
  assert isinstance(points, DaskDataFrame)
 
  transformations = _get_transformations_from_ngff_dict(f.attrs.asdict()["coordinateTransformations"])
@@ -57,7 +49,8 @@ def write_points(
  t = _get_transformations(points)
 
  points_groups = group.require_group(name)
- path = Path(points_groups._store.path) / points_groups.path / "points.parquet"
+ store = points_groups._store
+ new_path = os.path.join(store.path, points_groups.path, "points.parquet")
 
  # The following code iterates through all columns in the 'points' DataFrame. If the column's datatype is
  # 'category', it checks whether the categories of this column are known. If not, it explicitly converts the
@@ -70,7 +63,7 @@ def write_points(
  c = c.cat.as_known()
  points[column_name] = c
 
- points.to_parquet(path)
+ points.to_parquet(new_path, filesystem=getattr(store, "fs", None))
 
  attrs = format.attrs_to_dict(points.attrs)
  attrs["version"] = format.spatialdata_format_version