Validate element-type selection when reading Zarr stores

src/spatialdata/__main__.py

@@ -15,8 +15,8 @@
 
 @click.command(help="Peek inside the SpatialData .zarr dataset")
 @click.argument("path", default=False, type=str)
-@click.argument("selection", type=click.Choice(["images", "labels", "points", "shapes", "table"]), nargs=-1)
-def peek(path: str, selection: tuple[Literal["images", "labels", "points", "shapes", "table"]]) -> None:
+@click.argument("selection", type=click.Choice(["images", "labels", "points", "shapes", "tables"]), nargs=-1)
+def peek(path: str, selection: tuple[Literal["images", "labels", "points", "shapes", "tables"], ...]) -> None:
     """
     Peek inside the SpatialData .zarr dataset.
 
@@ -29,7 +29,7 @@ def peek(path: str, selection: tuple[Literal["images", "labels", "points", "shap
     path
         The path to the .zarr dataset to be inspected.
     selection
-        Optional, a list of keys (among images, labels, points, shapes, table) to load only a subset of the dataset.
+        Optional, a list of keys (among images, labels, points, shapes, tables) to load only a subset of the dataset.
         Example: `python -m spatialdata peek data.zarr images labels`
     """
     import spatialdata as sd
@@ -46,7 +46,7 @@ def peek(path: str, selection: tuple[Literal["images", "labels", "points", "shap
             f"Error: .zarr storage not found at {path}. Please specify a valid OME-NGFF spatial data (.zarr) file. "
             "Examples "
             '"python -m spatialdata peek data.zarr"'
-            '"python -m spatialdata peek https://remote/.../data.zarr labels table"'
+            '"python -m spatialdata peek https://remote/.../data.zarr labels tables"'
         )
 
 

src/spatialdata/_core/spatialdata.py

@@ -1960,7 +1960,7 @@ def tables(self, tables: dict[str, AnnData]) -> None:
     @staticmethod
     def read(
         file_path: str | Path | UPath,
-        selection: tuple[str] | None = None,
+        selection: tuple[Literal["images", "labels", "points", "shapes", "tables"], ...] | None = None,
         reconsolidate_metadata: bool = False,
     ) -> SpatialData:
         """
@@ -1977,7 +1977,8 @@ def read(
             The path or URL to the Zarr storage. To read from an already-open :class:`zarr.Group`, use
             :func:`spatialdata.read_zarr` instead.
         selection
-            The elements to read (images, labels, points, shapes, table). If None, all elements are read.
+            Tuple of element types to read: ``"images"``, ``"labels"``, ``"points"``, ``"shapes"``, or ``"tables"``.
+            If None or empty, all element types are read. Invalid values raise a :class:`ValueError`.
         reconsolidate_metadata
             If `True`, rewrite the consolidated metadata of the store before reading it. Use this when the consolidated
             metadata is corrupted or out of date, which otherwise leads to errors when reading the data. This requires

src/spatialdata/_io/io_zarr.py

@@ -133,7 +133,7 @@ def get_raster_format_for_read(
 
 def read_zarr(
     store: str | Path | UPath | zarr.Group,
-    selection: None | tuple[str] = None,
+    selection: tuple[Literal["images", "labels", "points", "shapes", "tables"], ...] | None = None,
     on_bad_files: Literal[BadFileHandleMethod.ERROR, BadFileHandleMethod.WARN] = BadFileHandleMethod.ERROR,
 ) -> SpatialData:
     """
@@ -150,8 +150,8 @@ def read_zarr(
         Path, URL, or zarr.Group to the zarr store (on-disk or remote).
 
     selection
-        List of elements to read from the zarr store (images, labels, points, shapes, tables). If None, all elements are
-        read.
+        Tuple of element types to read from the zarr store: ``"images"``, ``"labels"``, ``"points"``, ``"shapes"``,
+        or ``"tables"``. If None or empty, all element types are read. Invalid values raise a :class:`ValueError`.
 
     on_bad_files
         Specifies what to do upon encountering a bad file, e.g. corrupted, invalid or missing files.
@@ -173,6 +173,14 @@ def read_zarr(
     """
     from spatialdata._io._utils import _resolve_zarr_store
 
+    allowed_selection = {"images", "labels", "points", "shapes", "tables"}
+    if isinstance(selection, str) or (selection is not None and not set(selection).issubset(allowed_selection)):
+        raise ValueError(
+            f"Invalid selection: {selection!r}. Expected a tuple containing only {sorted(allowed_selection)}, "
+            "or None to read all element types."
+        )
+    selector: set[str] = allowed_selection if not selection else set(selection)
+
     resolved_store = _resolve_zarr_store(store)
     root_group = zarr.open_group(resolved_store, mode="r")
     # the following is the SpatialDataContainerFormat version
@@ -199,7 +207,6 @@ def read_zarr(
     shapes: dict[str, GeoDataFrame] = {}
     tables: dict[str, AnnData] = {}
 
-    selector = {"images", "labels", "points", "shapes", "tables"} if not selection else set(selection or [])
     logger.debug(f"Reading selection {selector}")
 
     # we could make this more readable. One can get lost when looking at this dict and iteration over the items

tests/io/test_read_zarr_selection.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pytest
+
+from spatialdata import SpatialData, read_zarr
+from spatialdata.models import Image2DModel
+
+
+@pytest.fixture
+def image_store(tmp_path: Path) -> Path:
+    path = tmp_path / "image.zarr"
+    image = Image2DModel.parse(np.zeros((1, 4, 4)), dims=("c", "y", "x"))
+    SpatialData(images={"image": image}).write(path)
+    return path
+
+
+@pytest.mark.parametrize("reader", [read_zarr, SpatialData.read])
+@pytest.mark.parametrize("selection", [("table",), ("imagez",), ("images", "table"), "images"])
+def test_invalid_selection(image_store: Path, reader: Any, selection: Any) -> None:
+    with pytest.raises(ValueError, match="Invalid selection"):
+        reader(image_store, selection=selection)
+
+
+@pytest.mark.parametrize("reader", [read_zarr, SpatialData.read])
+@pytest.mark.parametrize("selection", [None, (), ("images",), ("images", "tables"), ("images", "images")])
+def test_valid_selection_keeps_images(image_store: Path, reader: Any, selection: Any) -> None:
+    result = reader(image_store, selection=selection)
+    assert list(result.images) == ["image"]
+    np.testing.assert_array_equal(result.images["image"].values, np.zeros((1, 4, 4)))
+
+
+@pytest.mark.parametrize("selection", [("labels",), ("points",), ("shapes",), ("tables",)])
+def test_valid_selection_of_absent_type(image_store: Path, selection: Any) -> None:
+    result = read_zarr(image_store, selection=selection)
+    assert list(result.gen_elements()) == []

tests/test_cli.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+from anndata import AnnData
+from click.testing import CliRunner
+
+from spatialdata import SpatialData
+from spatialdata.__main__ import cli
+from spatialdata.models import Image2DModel
+
+
+@pytest.mark.parametrize("selection", [[], ["tables"], ["images", "tables"]])
+def test_peek_tables(tmp_path: Path, selection: list[str]) -> None:
+    path = tmp_path / "data.zarr"
+    image = Image2DModel.parse(np.zeros((1, 4, 4)), dims=("c", "y", "x"))
+    SpatialData(images={"test_image": image}, tables={"test_table": AnnData(np.ones((2, 3)))}).write(path)
+
+    result = CliRunner().invoke(cli, ["peek", str(path), *selection])
+
+    assert result.exit_code == 0, result.output
+    loaded_elements = result.output.split("with the following elements in the Zarr store")[0]
+    assert "test_table" in loaded_elements
+    assert ("test_image" in loaded_elements) == (not selection or "images" in selection)