diff --git a/src/spatialdata/__main__.py b/src/spatialdata/__main__.py index 6d15de398..5217c655e 100644 --- a/src/spatialdata/__main__.py +++ b/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"' ) diff --git a/src/spatialdata/_core/spatialdata.py b/src/spatialdata/_core/spatialdata.py index 89f4a4920..6a40a50c4 100644 --- a/src/spatialdata/_core/spatialdata.py +++ b/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 diff --git a/src/spatialdata/_io/io_zarr.py b/src/spatialdata/_io/io_zarr.py index 186d7c974..31c8f8f38 100644 --- a/src/spatialdata/_io/io_zarr.py +++ b/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 diff --git a/tests/io/test_read_zarr_selection.py b/tests/io/test_read_zarr_selection.py new file mode 100644 index 000000000..dc5354d11 --- /dev/null +++ b/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()) == [] diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 000000000..8ecb302e6 --- /dev/null +++ b/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)