from __future__ import annotations

import json

import os

import tempfile

from collections.abc import Callable

from importlib.metadata import version

from pathlib import Path

from typing import Any, Literal

import dask.array as da

import dask.dataframe as dd

import numpy as np

import pandas as pd

import pyarrow.parquet as pq

import pytest

import zarr

from anndata import AnnData

from anndata.io import read_elem

from numpy.random import default_rng

from packaging.version import Version

from pandas.testing import assert_series_equal

from shapely import MultiPolygon, Polygon

from upath import UPath

from xarray import DataArray

from zarr.errors import GroupNotFoundError

import spatialdata.config

from spatialdata import SpatialData, deepcopy, read_zarr

from spatialdata._core.validation import ValidationError

from spatialdata._io._utils import _are_directories_identical, get_dask_backing_files

from spatialdata._io.format import (

CurrentSpatialDataContainerFormat,

SpatialDataContainerFormats,

SpatialDataContainerFormatType,

SpatialDataContainerFormatV01,

)

from spatialdata._io.io_raster import write_image

from spatialdata.datasets import blobs

from spatialdata.models import Image2DModel

from spatialdata.models._utils import get_channel_names

from spatialdata.testing import assert_spatial_data_objects_are_identical

from spatialdata.transformations.operations import (

get_transformation,

set_transformation,

)

from spatialdata.transformations.transformations import Identity, Scale

from tests.conftest import (

_get_images,

_get_labels,

_get_points,

_get_shapes,

_get_table,

_get_tables,

)

RNG = default_rng(0)

SDATA_FORMATS = list(SpatialDataContainerFormats.values())

@pytest.mark.filterwarnings("ignore:SpatialData is not stored in the most current format:UserWarning")

@pytest.mark.parametrize("sdata_container_format", SDATA_FORMATS)

class TestReadWrite:

def test_images(

self,

tmp_path: str,

images: SpatialData,

sdata_container_format: SpatialDataContainerFormatType,

) -> None:

tmpdir = Path(tmp_path) / "tmp.zarr"

# ensures that we are inplicitly testing the read and write of channel names

assert get_channel_names(images["image2d"]) == ["r", "g", "b"]

assert get_channel_names(images["image2d_multiscale"]) == ["r", "g", "b"]

images.write(tmpdir, sdata_formats=sdata_container_format)

sdata = SpatialData.read(tmpdir)

assert_spatial_data_objects_are_identical(images, sdata)

def test_labels(

self,

tmp_path: str,

labels: SpatialData,

sdata_container_format: SpatialDataContainerFormatType,

) -> None:

tmpdir = Path(tmp_path) / "tmp.zarr"

labels.write(tmpdir, sdata_formats=sdata_container_format)

sdata = SpatialData.read(tmpdir)

assert_spatial_data_objects_are_identical(labels, sdata)

@pytest.mark.parametrize("geometry_encoding", ["WKB", "geoarrow"])

def test_shapes(

self,

tmp_path: str,

shapes: SpatialData,

sdata_container_format: SpatialDataContainerFormatType,

geometry_encoding: Literal["WKB", "geoarrow"],

) -> None:

tmpdir = Path(tmp_path) / "tmp.zarr"

# check the index is correctly written and then read

shapes["circles"].index = np.arange(1, len(shapes["circles"]) + 1)

# add a mixed Polygon + MultiPolygon element

shapes["mixed"] = pd.concat([shapes["poly"], shapes["multipoly"]])

shapes.write(tmpdir, sdata_formats=sdata_container_format, shapes_geometry_encoding=geometry_encoding)

sdata = SpatialData.read(tmpdir)

if geometry_encoding == "WKB":

assert_spatial_data_objects_are_identical(shapes, sdata)

else:

# convert each Polygon to a MultiPolygon

mixed_multipolygon = shapes["mixed"].assign(

geometry=lambda df: df.geometry.apply(lambda g: MultiPolygon([g]) if isinstance(g, Polygon) else g)

)

assert sdata["mixed"].equals(mixed_multipolygon)

assert not sdata["mixed"].equals(shapes["mixed"])

del shapes["mixed"]

del sdata["mixed"]

assert_spatial_data_objects_are_identical(shapes, sdata)

@pytest.mark.parametrize("geometry_encoding", ["WKB", "geoarrow"])

def test_shapes_geometry_encoding_write_element(

self,

tmp_path: str,

shapes: SpatialData,

sdata_container_format: SpatialDataContainerFormatType,

geometry_encoding: Literal["WKB", "geoarrow"],

) -> None:

"""Test shapes geometry encoding with write_element() and global settings."""

tmpdir = Path(tmp_path) / "tmp.zarr"

# First write an empty SpatialData to create the zarr store

empty_sdata = SpatialData()

empty_sdata.write(tmpdir, sdata_formats=sdata_container_format)

shapes["mixed"] = pd.concat([shapes["poly"], shapes["multipoly"]])

# Add shapes to the empty sdata

for shape_name in shapes.shapes:

empty_sdata[shape_name] = shapes[shape_name]

# Store original setting and set global encoding

original_encoding = spatialdata.config.settings.shapes_geometry_encoding

try:

spatialdata.config.settings.shapes_geometry_encoding = geometry_encoding

# Write each shape element - should use global setting

for shape_name in shapes.shapes:

empty_sdata.write_element(shape_name, sdata_formats=sdata_container_format)

# Verify the encoding metadata in the parquet file

parquet_file = tmpdir / "shapes" / shape_name / "shapes.parquet"

with pq.ParquetFile(parquet_file) as pf:

md = pf.metadata

d = json.loads(md.metadata[b"geo"].decode("utf-8"))

found_encoding = d["columns"]["geometry"]["encoding"]

if geometry_encoding == "WKB":

expected_encoding = "WKB"

elif shape_name == "circles":

expected_encoding = "point"

elif shape_name == "poly":

expected_encoding = "polygon"

elif shape_name in ["multipoly", "mixed"]:

expected_encoding = "multipolygon"

else:

raise ValueError(

f"Uncovered case for shape_name: {shape_name}, found encoding: {found_encoding}."

)

assert found_encoding == expected_encoding

finally:

spatialdata.config.settings.shapes_geometry_encoding = original_encoding

def test_points(

self,

tmp_path: str,

points: SpatialData,

sdata_container_format: SpatialDataContainerFormatType,

) -> None:

tmpdir = Path(tmp_path) / "tmp.zarr"

# check the index is correctly written and then read

new_index = dd.from_array(np.arange(1, len(points["points_0"]) + 1))

points["points_0"] = points["points_0"].set_index(new_index)

points.write(tmpdir, sdata_formats=sdata_container_format)

sdata = SpatialData.read(tmpdir)

assert_spatial_data_objects_are_identical(points, sdata)

def _test_table(

self,

tmp_path: str,

table: SpatialData,

sdata_container_format: SpatialDataContainerFormatType,

) -> None:

tmpdir = Path(tmp_path) / "tmp.zarr"

table.write(tmpdir, sdata_formats=sdata_container_format)

sdata = SpatialData.read(tmpdir)

assert_spatial_data_objects_are_identical(table, sdata)

def test_single_table_single_annotation(

self,

tmp_path: str,

table_single_annotation: SpatialData,

sdata_container_format: SpatialDataContainerFormatType,

) -> None:

self._test_table(

tmp_path,

table_single_annotation,

sdata_container_format=sdata_container_format,

)

def test_single_table_multiple_annotations(

self,

tmp_path: str,

table_multiple_annotations: SpatialData,

sdata_container_format: SpatialDataContainerFormatType,

) -> None:

self._test_table(

tmp_path,

table_multiple_annotations,

sdata_container_format=sdata_container_format,

)

def test_multiple_tables(

self,

tmp_path: str,

tables: list[AnnData],

sdata_container_format: SpatialDataContainerFormatType,

) -> None:

sdata_tables = SpatialData(tables={str(i): tables[i] for i in range(len(tables))})

self._test_table(tmp_path, sdata_tables, sdata_container_format=sdata_container_format)

def test_roundtrip(

self,

tmp_path: str,

sdata: SpatialData,

sdata_container_format: SpatialDataContainerFormatType,

) -> None:

tmpdir = Path(tmp_path) / "tmp.zarr"

sdata.write(tmpdir, sdata_formats=sdata_container_format)

sdata2 = SpatialData.read(tmpdir)

tmpdir2 = Path(tmp_path) / "tmp2.zarr"

sdata2.write(tmpdir2, sdata_formats=sdata_container_format)

_are_directories_identical(tmpdir, tmpdir2, exclude_regexp="[1-9][0-9]*.*")

def test_compression_roundtrip(

self,

tmp_path: str,

full_sdata: SpatialData,

sdata_container_format: SpatialDataContainerFormatType,

):

tmpdir = Path(tmp_path) / "tmp.zarr"

with pytest.raises(TypeError, match="Expected a dictionary with as"):

full_sdata.write(tmpdir, raster_compressor="faulty", sdata_formats=sdata_container_format)

with pytest.raises(ValueError, match="Expected a dictionary with a single"):

full_sdata.write(

tmpdir, raster_compressor={"zstd": 8, "other_item": 4}, sdata_formats=sdata_container_format

)

with pytest.raises(ValueError, match="Compression must either"):

full_sdata.write(tmpdir, raster_compressor={"faulty": 8}, sdata_formats=sdata_container_format)

with pytest.raises(ValueError, match="Compression must either"):

full_sdata.write(

tmpdir, raster_compressor={"The compression level": 10}, sdata_formats=sdata_container_format

)

full_sdata.write(tmpdir, raster_compressor={"zstd": 8}, sdata_formats=sdata_container_format)

# sourcery skip: no-loop-in-tests

for element in ["image2d", "image2d_multiscale", "labels2d", "labels2d_multiscale"]:

element_type = "images" if element.startswith("image") else "labels"

arr = zarr.open_group(tmpdir / element_type, mode="r")[element]["s0"]

compressor = arr.compressors[0]

if sdata_container_format.zarr_format == 2:

assert compressor.cname == "zstd"

assert compressor.clevel == 8

elif sdata_container_format.zarr_format == 3:

from zarr.codecs.zstd import ZstdCodec

assert isinstance(compressor, ZstdCodec)

assert compressor.level == 8

@pytest.mark.parametrize("raster_compressor", [{"lz4": 3}, {"zstd": 7}])

@pytest.mark.parametrize("element", [("images", "image2d"), ("labels", "labels2d")])

def test_write_element_compression(

self,

tmp_path: str,

full_sdata: SpatialData,

raster_compressor: dict[Literal["lz4", "zstd"], int],

element: str,

sdata_container_format: SpatialDataContainerFormatType,

):

tmpdir = Path(tmp_path) / "compression.zarr"

sdata = SpatialData()

sdata.write(tmpdir, sdata_formats=sdata_container_format)

sdata["element"] = full_sdata[element[1]]

sdata.write_element("element", raster_compressor=raster_compressor, sdata_formats=sdata_container_format)

arr = zarr.open_group(tmpdir / element[0], mode="r")["element"]["s0"]

compression = arr.compressors[0]

if sdata_container_format.zarr_format == 2:

assert compression.cname == list(raster_compressor.keys())[0]

assert compression.clevel == list(raster_compressor.values())[0]

elif sdata_container_format.zarr_format == 3:

from zarr.codecs import BloscCodec, ZstdCodec

compressor_name = list(raster_compressor.keys())[0]

compressor_level = list(raster_compressor.values())[0]

if compressor_name == "zstd":

assert isinstance(compression, ZstdCodec)

assert compression.level == compressor_level

elif compressor_name == "lz4":

assert isinstance(compression, BloscCodec)

assert compression.clevel == compressor_level

def test_incremental_io_list_of_elements(

self,

shapes: SpatialData,

sdata_container_format: SpatialDataContainerFormatType,

) -> None:

with tempfile.TemporaryDirectory() as tmpdir:

f = os.path.join(tmpdir, "data.zarr")

shapes.write(f, sdata_formats=sdata_container_format)

new_shapes0 = deepcopy(shapes["circles"])

new_shapes1 = deepcopy(shapes["poly"])

shapes["new_shapes0"] = new_shapes0

shapes["new_shapes1"] = new_shapes1

assert "shapes/new_shapes0" not in shapes.elements_paths_on_disk()

assert "shapes/new_shapes1" not in shapes.elements_paths_on_disk()

shapes.write_element(["new_shapes0", "new_shapes1"], sdata_formats=sdata_container_format)

assert "shapes/new_shapes0" in shapes.elements_paths_on_disk()

assert "shapes/new_shapes1" in shapes.elements_paths_on_disk()

shapes.delete_element_from_disk(["new_shapes0", "new_shapes1"])

assert "shapes/new_shapes0" not in shapes.elements_paths_on_disk()

assert "shapes/new_shapes1" not in shapes.elements_paths_on_disk()

@staticmethod

def _workaround1_non_dask_backed(

sdata: SpatialData,

name: str,

new_name: str,

sdata_container_format: SpatialDataContainerFormatType = CurrentSpatialDataContainerFormat(),

) -> None:

# a. write a backup copy of the data

sdata[new_name] = sdata[name]

sdata.write_element(new_name, sdata_formats=sdata_container_format)

# b. rewrite the original data

sdata.delete_element_from_disk(name)

sdata.write_element(name, sdata_formats=sdata_container_format)

# c. remove the backup copy

del sdata[new_name]

sdata.delete_element_from_disk(new_name)

@staticmethod

def _workaround1_dask_backed(

sdata: SpatialData,

name: str,

new_name: str,

sdata_container_format: SpatialDataContainerFormatType = CurrentSpatialDataContainerFormat(),

) -> None:

# a. write a backup copy of the data

sdata[new_name] = sdata[name]

sdata.write_element(new_name, sdata_formats=sdata_container_format)

# a2. remove the in-memory copy from the SpatialData object (note,

# at this point the backup copy still exists on-disk)

del sdata[new_name]

del sdata[name]

# a3 load the backup copy into memory

sdata_copy = read_zarr(sdata.path)

# b1. rewrite the original data

sdata.delete_element_from_disk(name)

sdata[name] = sdata_copy[new_name]

sdata.write_element(name, sdata_formats=sdata_container_format)

# b2. reload the new data into memory (because it has been written but in-memory it still points

# from the backup location)

sdata = read_zarr(sdata.path)

# c. remove the backup copy

del sdata[new_name]

sdata.delete_element_from_disk(new_name)

@pytest.mark.parametrize("dask_backed", [True, False])

@pytest.mark.parametrize("workaround", [1, 2])

def test_incremental_io_on_disk(

self,

tmp_path: str,

full_sdata: SpatialData,

dask_backed: bool,

workaround: int,

sdata_container_format: SpatialDataContainerFormatType,

) -> None:

"""

This tests shows workaround on how to rewrite existing data on disk.

The user is recommended to study them, understand the implications and eventually adapt them to their use case.

We are working on simpler workarounds and on a more robust solution to this problem, but unfortunately it is not

yet available.

In particular the complex "dask-backed" case for workaround 1 could be simplified once

"""

_elements = [

"image2d",

"image3d_multiscale_xarray",

"labels2d",

"labels3d_multiscale_xarray",

"points_0",

"multipoly",

"table",

]

# Reduce to only the elements under test so the fixture deepcopy stays small.

full_sdata = full_sdata.subset(_elements)

tmpdir = Path(tmp_path) / "incremental_io.zarr"

sdata = SpatialData()

sdata.write(tmpdir, sdata_formats=sdata_container_format)

for name in _elements:

sdata[name] = full_sdata[name]

sdata.write_element(name, sdata_formats=sdata_container_format)

if dask_backed:

# this forces the element to write to be dask-backed from disk. In this case, overwriting the data is

# more laborious because we are writing the data to the same location that defines the data!

sdata = read_zarr(sdata.path)

with pytest.raises(

ValueError,

match="The Zarr store already exists. Use `overwrite=True` to try overwriting the store.",

):

sdata.write_element(name, sdata_formats=sdata_container_format)

match = (

"Details: the target path contains one or more files that Dask use for backing elements in the "

"SpatialData object"

if dask_backed

and name

in [

"image2d",

"labels2d",

"image3d_multiscale_xarray",

"labels3d_multiscale_xarray",

"points_0",

]

else "Details: the target path in which to save an element is a subfolder of the current Zarr store."

)

with pytest.raises(

ValueError,

match=match,

):

sdata.write_element(name, overwrite=True, sdata_formats=sdata_container_format)

if workaround == 1:

new_name = f"{name}_new_place"

# workaround 1, mostly safe (untested for Windows platform, network drives, multi-threaded

# setups, ...). If the scenario matches your use case, please use with caution.

if not dask_backed: # easier case

self._workaround1_non_dask_backed(

sdata=sdata,

name=name,

new_name=new_name,

sdata_container_format=sdata_container_format,

)

else: # dask-backed case, more complex

self._workaround1_dask_backed(

sdata=sdata,

name=name,

new_name=new_name,

sdata_container_format=sdata_container_format,

)

elif workaround == 2:

# workaround 2, unsafe but sometimes acceptable depending on the user's workflow.

# this works only if the data is not dask-backed, otherwise an exception will be raised because the code

# would be trying to delete the data that the Dask object is pointing to!

if not dask_backed:

# a. rewrite the original data (risky!)

sdata.delete_element_from_disk(name)

sdata.write_element(name, sdata_formats=sdata_container_format)

def test_io_and_lazy_loading_points(self, points, sdata_container_format: SpatialDataContainerFormatType):

with tempfile.TemporaryDirectory() as td:

f = os.path.join(td, "data.zarr")

points.write(f, sdata_formats=sdata_container_format)

assert len(get_dask_backing_files(points)) == 0

sdata2 = SpatialData.read(f)

assert len(get_dask_backing_files(sdata2)) > 0

def test_io_and_lazy_loading_raster(self, images, labels, sdata_container_format: SpatialDataContainerFormatType):

sdatas = {"images": images, "labels": labels}

for k, sdata in sdatas.items():

d = getattr(sdata, k)

elem_name = list(d.keys())[0]

with tempfile.TemporaryDirectory() as td:

f = os.path.join(td, "data.zarr")

dask0 = sdata[elem_name].data

sdata.write(f, sdata_formats=sdata_container_format)

assert all("from-zarr" not in key for key in dask0.dask.layers)

assert len(get_dask_backing_files(sdata)) == 0

sdata2 = SpatialData.read(f)

dask1 = sdata2[elem_name].data

assert any("from-zarr" in key for key in dask1.dask.layers)

assert len(get_dask_backing_files(sdata2)) > 0

def test_replace_transformation_on_disk_raster(

self, images, labels, sdata_container_format: SpatialDataContainerFormatType

):

sdatas = {"images": images, "labels": labels}

for k, sdata in sdatas.items():

d = getattr(sdata, k)

# unlike the non-raster case, we are testing all the elements (2d and 3d, multiscale and not)

for elem_name in d:

kwargs = {k: {elem_name: d[elem_name]}}

single_sdata = SpatialData(**kwargs)

with tempfile.TemporaryDirectory() as td:

f = os.path.join(td, "data.zarr")

single_sdata.write(f, sdata_formats=sdata_container_format)

t0 = get_transformation(SpatialData.read(f)[elem_name])

assert isinstance(t0, Identity)

set_transformation(

single_sdata[elem_name],

Scale([2.0], axes=("x",)),

write_to_sdata=single_sdata,

)

t1 = get_transformation(SpatialData.read(f)[elem_name])

assert isinstance(t1, Scale)

def test_replace_transformation_on_disk_non_raster(

self, shapes, points, sdata_container_format: SpatialDataContainerFormatType

):

sdatas = {"shapes": shapes, "points": points}

for k, sdata in sdatas.items():

d = sdata.__getattribute__(k)

elem_name = list(d.keys())[0]

with tempfile.TemporaryDirectory() as td:

f = os.path.join(td, "data.zarr")

sdata.write(f, sdata_formats=sdata_container_format)

t0 = get_transformation(SpatialData.read(f).__getattribute__(k)[elem_name])

assert isinstance(t0, Identity)

set_transformation(sdata[elem_name], Scale([2.0], axes=("x",)), write_to_sdata=sdata)

t1 = get_transformation(SpatialData.read(f)[elem_name])

assert isinstance(t1, Scale)

def test_write_overwrite_fails_when_no_zarr_store(

self, full_sdata, sdata_container_format: SpatialDataContainerFormatType

):

with tempfile.TemporaryDirectory() as tmpdir:

f = Path(tmpdir) / "data.zarr"

f.mkdir()

old_data = SpatialData()

with pytest.raises(ValueError, match="The target file path specified already exists"):

old_data.write(f, sdata_formats=sdata_container_format)

with pytest.raises(ValueError, match="The target file path specified already exists"):

full_sdata.write(f, overwrite=True, sdata_formats=sdata_container_format)

def test_overwrite_fails_when_no_zarr_store_but_dask_backed_data(

self,

full_sdata,

points,

images,

labels,

sdata_container_format: SpatialDataContainerFormatType,

):

sdatas = {"images": images, "labels": labels, "points": points}

elements = {"images": "image2d", "labels": "labels2d", "points": "points_0"}

for k, sdata in sdatas.items():

element = elements[k]

with tempfile.TemporaryDirectory() as tmpdir:

f = os.path.join(tmpdir, "data.zarr")

sdata.write(f, sdata_formats=sdata_container_format)

# now we have a sdata with dask-backed elements

sdata2 = SpatialData.read(f)

p = sdata2[element]

full_sdata[element] = p

with pytest.raises(

ValueError,

match="The Zarr store already exists. Use `overwrite=True` to try overwriting the store.",

):

full_sdata.write(f, sdata_formats=sdata_container_format)

with pytest.raises(

ValueError,

match=r"Details: the target path contains one or more files that Dask use for "

"backing elements in the SpatialData object",

):

full_sdata.write(f, overwrite=True, sdata_formats=sdata_container_format)

def test_overwrite_fails_when_zarr_store_present(self, sdata_container_format: SpatialDataContainerFormatType):

# addressing https://github.com/scverse/spatialdata/issues/137

with tempfile.TemporaryDirectory() as tmpdir:

f = os.path.join(tmpdir, "data.zarr")

# An empty store is enough to trigger both exceptions.

sdata = SpatialData()

sdata.write(f, sdata_formats=sdata_container_format)

with pytest.raises(

ValueError,

match="The Zarr store already exists. Use `overwrite=True` to try overwriting the store.",

):

sdata.write(f, sdata_formats=sdata_container_format)

with pytest.raises(

ValueError,

match=r"Details: the target path either contains, coincides or is contained in the current Zarr store",

):

sdata.write(f, overwrite=True, sdata_formats=sdata_container_format)

# support for overwriting backed sdata has been temporarily removed

# with tempfile.TemporaryDirectory() as tmpdir:

# f = os.path.join(tmpdir, "data.zarr")

# full_sdata.write(f)

# full_sdata.write(f, overwrite=True)

# print(full_sdata)

#

# sdata2 = SpatialData(

# images=full_sdata.images,

# labels=full_sdata.labels,

# points=full_sdata.points,

# shapes=full_sdata.shapes,

# table=full_sdata.table,

# )

# sdata2.write(f, overwrite=True)

def test_overwrite_fails_onto_non_zarr_file(

self, full_sdata, sdata_container_format: SpatialDataContainerFormatType

):

ERROR_MESSAGE = (

"The target file path specified already exists, and it has been detected to not be a Zarr store."

)

with tempfile.TemporaryDirectory() as tmpdir:

f0 = os.path.join(tmpdir, "test.txt")

with open(f0, "w"):

with pytest.raises(

ValueError,

match=ERROR_MESSAGE,

):

full_sdata.write(f0, sdata_formats=sdata_container_format)

with pytest.raises(

ValueError,

match=ERROR_MESSAGE,

):

full_sdata.write(f0, overwrite=True, sdata_formats=sdata_container_format)

f1 = os.path.join(tmpdir, "test.zarr")

os.mkdir(f1)

with pytest.raises(ValueError, match=ERROR_MESSAGE):

full_sdata.write(f1, sdata_formats=sdata_container_format)

with pytest.raises(ValueError, match=ERROR_MESSAGE):

full_sdata.write(f1, overwrite=True, sdata_formats=sdata_container_format)

def test_incremental_io_in_memory(

full_sdata: SpatialData,

) -> None:

sdata = full_sdata

for k, v in _get_images().items():

sdata.images[f"additional_{k}"] = v

with pytest.raises(KeyError, match="Key `table` is not unique"):

sdata["table"] = v

for k, v in _get_labels().items():

sdata.labels[f"additional_{k}"] = v

with pytest.raises(KeyError, match="Key `table` is not unique"):

sdata["table"] = v

for k, v in _get_shapes().items():

sdata.shapes[f"additional_{k}"] = v

with pytest.raises(KeyError, match="Key `table` is not unique"):

sdata["table"] = v

for k, v in _get_points().items():

sdata.points[f"additional_{k}"] = v

with pytest.raises(KeyError, match="Key `table` is not unique"):

sdata["table"] = v

for k, v in _get_tables(region="labels2d").items():

sdata.tables[f"additional_{k}"] = v

with pytest.raises(KeyError, match="Key `poly` is not unique"):

sdata["poly"] = v

def test_table_group_keeps_anndata_encoding_metadata(tmp_path: str, table_single_annotation: SpatialData) -> None:

# https://github.com/scverse/spatialdata/issues/1183

# Writing the spatialdata attributes on the table group must not erase the

# `encoding-type`/`encoding-version` metadata that anndata writes on the same

# group; anndata-level readers (read_elem, read_lazy) dispatch on it.

tmpdir = Path(tmp_path) / "tmp.zarr"

table_single_annotation.write(tmpdir)

on_disk = json.loads((tmpdir / "tables" / "table" / "zarr.json").read_text())["attributes"]

assert on_disk["encoding-type"] == "anndata"

assert on_disk["encoding-version"] == "0.1.0"

# the spatialdata attributes must be present alongside, not instead

assert on_disk["spatialdata-encoding-type"] == "ngff:regions_table"

table_group = zarr.open_group(tmpdir, mode="r")["tables/table"]

assert isinstance(read_elem(table_group), AnnData)

def test_bug_rechunking_after_queried_raster():

# https://github.com/scverse/spatialdata-io/issues/117

##

single_scale = Image2DModel.parse(RNG.random((100, 10, 10)), chunks=(5, 5, 5))

multi_scale = Image2DModel.parse(RNG.random((100, 10, 10)), scale_factors=[2, 2], chunks=(5, 5, 5))

images = {"single_scale": single_scale, "multi_scale": multi_scale}

sdata = SpatialData(images=images)

queried = sdata.query.bounding_box(

axes=("x", "y"),

min_coordinate=[2, 5],

max_coordinate=[12, 12],

target_coordinate_system="global",

)

with tempfile.TemporaryDirectory() as tmpdir:

f = os.path.join(tmpdir, "data.zarr")

queried.write(f)

def test_is_regular_dask_chunk_grid() -> None:

from spatialdata._io.io_raster import _is_regular_dask_chunk_grid

# Single chunk per axis → continue branch, overall True

assert _is_regular_dask_chunk_grid([(4,)]) is True

# Empty axis → continue branch, overall True

assert _is_regular_dask_chunk_grid([()]) is True

# Non-uniform interior chunks → first return False

assert _is_regular_dask_chunk_grid([(4, 4, 3, 4)]) is False

# Last chunk larger than first → second return False

assert _is_regular_dask_chunk_grid([(4, 4, 4, 5)]) is False

# All chunks equal → True

assert _is_regular_dask_chunk_grid([(4, 4, 4, 4)]) is True

# Last chunk smaller than first → True

assert _is_regular_dask_chunk_grid([(4, 4, 4, 1)]) is True

# Empty grid (no axes) → True

assert _is_regular_dask_chunk_grid([]) is True

# Multi-axis: all axes regular → True

assert _is_regular_dask_chunk_grid([(4, 4, 4, 1), (3, 3, 2)]) is True

# Multi-axis: one axis irregular → False

assert _is_regular_dask_chunk_grid([(4, 4, 4, 1), (4, 4, 3, 4)]) is False

def test_write_irregular_dask_chunks_without_explicit_storage_options(tmp_path: Path) -> None:

data = da.from_array(RNG.random((3, 800, 1000)), chunks=((3,), (300, 200, 300), (512, 488)))

image = Image2DModel.parse(data, dims=("c", "y", "x"))

sdata = SpatialData(images={"image": image})

path = tmp_path / "data.zarr"

with pytest.warns(UserWarning, match="irregular chunk sizes"):

sdata.write(path)

sdata_back = read_zarr(path)

assert sdata_back["image"].chunks == ((3,), (300, 300, 200), (512, 488))

def test_write_image_normalizes_explicit_regular_dask_chunk_grid(tmp_path: Path) -> None:

data = da.from_array(RNG.random((3, 800, 1000)), chunks=((3,), (300, 300, 200), (512, 488)))

image = Image2DModel.parse(data, dims=("c", "y", "x"))

group = zarr.open_group(tmp_path / "image.zarr", mode="w")

write_image(image, group, "image", storage_options={"chunks": image.data.chunks})

assert group["s0"].chunks == (3, 300, 512)

def test_write_image_rejects_explicit_irregular_dask_chunk_grid(tmp_path: Path) -> None:

data = da.from_array(RNG.random((3, 800, 1000)), chunks=((3,), (300, 200, 300), (512, 488)))

image = Image2DModel.parse(data, dims=("c", "y", "x"))

group = zarr.open_group(tmp_path / "image.zarr", mode="w")

with pytest.raises(

ValueError,

match='storage_options\\["chunks"\\] must resolve to a Zarr chunk shape or a regular Dask chunk grid',

):

write_image(image, group, "image", storage_options={"chunks": image.data.chunks})

def test_write_image_normalizes_explicit_zarr_chunk_grid(tmp_path: Path) -> None:

data = da.from_array(RNG.random((3, 800, 1000)), chunks=((3,), (300, 200, 300), (512, 488)))

image = Image2DModel.parse(data, dims=("c", "y", "x"))

group = zarr.open_group(tmp_path / "image.zarr", mode="w")

zarr_chunks = (3, 100, 512) # ome zarr rechunks when writing

write_image(image, group, "image", storage_options={"chunks": zarr_chunks})

assert group["s0"].chunks == (3, 100, 512)

def test_write_image_rejects_string(tmp_path: Path) -> None:

data = da.from_array(RNG.random((3, 800, 1000)), chunks=((3,), (300, 300, 200), (512, 488)))

image = Image2DModel.parse(data, dims=("c", "y", "x"))

group = zarr.open_group(tmp_path / "image.zarr", mode="w")

with pytest.raises(

ValueError,

match='storage_options\\["chunks"\\] must resolve to a Zarr chunk shape or a regular Dask chunk grid',

):

write_image(image, group, "image", storage_options={"chunks": "auto"})

def test_write_image_rejects_empty_string(tmp_path: Path) -> None:

data = da.from_array(RNG.random((3, 800, 1000)), chunks=((3,), (300, 300, 200), (512, 488)))

image = Image2DModel.parse(data, dims=("c", "y", "x"))

group = zarr.open_group(tmp_path / "image.zarr", mode="w")

with pytest.raises(

ValueError,

match='storage_options\\["chunks"\\] must resolve to a Zarr chunk shape or a regular Dask chunk grid',

):

write_image(image, group, "image", storage_options={"chunks": ""})

def test_write_image_rejects_byte_string(tmp_path: Path) -> None:

data = da.from_array(RNG.random((3, 800, 1000)), chunks=((3,), (300, 300, 200), (512, 488)))

image = Image2DModel.parse(data, dims=("c", "y", "x"))

group = zarr.open_group(tmp_path / "image.zarr", mode="w")

with pytest.raises(

ValueError,

match='storage_options\\["chunks"\\] must resolve to a Zarr chunk shape or a regular Dask chunk grid',

):

write_image(image, group, "image", storage_options={"chunks": b"auto"})

def test_single_scale_image_roundtrip_stays_dataarray(tmp_path: Path) -> None:

image = Image2DModel.parse(RNG.random((3, 64, 64)), dims=("c", "y", "x"))

sdata = SpatialData(images={"image": image})

path = tmp_path / "data.zarr"

sdata.write(path)

sdata_back = read_zarr(path)

assert isinstance(sdata_back["image"], DataArray)

image_group = zarr.open_group(path / "images" / "image", mode="r")

assert list(image_group.keys()) == ["s0"]

@pytest.mark.filterwarnings("ignore:SpatialData is not stored in the most current format:UserWarning")

@pytest.mark.parametrize("sdata_container_format", SDATA_FORMATS)

def test_self_contained(full_sdata: SpatialData, sdata_container_format: SpatialDataContainerFormatType) -> None:

# image2d/labels2d/points_0 are used explicitly in the combined-element block below;

# circles covers the "else: assert self_contained" branch (non-dask-backed elements).

full_sdata = full_sdata.subset(["image2d", "labels2d", "points_0", "circles"])

# data only in-memory, so the SpatialData object and all its elements are self-contained

assert full_sdata.is_self_contained()

description = full_sdata.elements_are_self_contained()

assert all(description.values())

with tempfile.TemporaryDirectory() as tmpdir:

# data saved to disk, it's self contained

f = os.path.join(tmpdir, "data.zarr")

full_sdata.write(f, sdata_formats=sdata_container_format)

full_sdata.is_self_contained()

# we read the data, so it's self-contained

sdata2 = SpatialData.read(f)

assert sdata2.is_self_contained()

# we save the data to a new location, so it's not self-contained anymore

f2 = os.path.join(tmpdir, "data2.zarr")

sdata2.write(f2, sdata_formats=sdata_container_format)

assert not sdata2.is_self_contained()

# because of the images, labels and points

description = sdata2.elements_are_self_contained()

for element_name, self_contained in description.items():

if any(element_name.startswith(prefix) for prefix in ["image", "labels", "points"]):

assert not self_contained

else:

assert self_contained

# but if we read it again, it's self-contained

sdata3 = SpatialData.read(f2)

assert sdata3.is_self_contained()

# or if we do some more targeted manipulation

sdata2.path = Path(f)

assert sdata2.is_self_contained()

# finally, an example of a non-self-contained element which is not depending on files external to the Zarr store

# here we create an element which combines lazily 3 elements of the Zarr store; it will be a nonsense element,

# but good for testing

v = sdata2["points_0"]["x"].loc[0].values

v.compute_chunk_sizes()

combined = (

v

+ sdata2["labels2d"].expand_dims("c", 1).transpose("c", "y", "x")

+ sdata2["image2d"].sel(c="r").expand_dims("c", 1)

+ v

)

combined = Image2DModel.parse(combined)

assert len(get_dask_backing_files(combined)) == 3

sdata2["combined"] = combined

assert not sdata2.is_self_contained()

description = sdata2.elements_are_self_contained()

assert description["combined"] is False

assert all(description[element_name] for element_name in description if element_name != "combined")

@pytest.mark.filterwarnings("ignore:SpatialData is not stored in the most current format:UserWarning")

@pytest.mark.parametrize("sdata_container_format", SDATA_FORMATS)

def test_symmetric_difference_with_zarr_store(

full_sdata: SpatialData, sdata_container_format: SpatialDataContainerFormatType

) -> None:

with tempfile.TemporaryDirectory() as tmpdir:

f = os.path.join(tmpdir, "data.zarr")

full_sdata.write(f, sdata_formats=sdata_container_format)

# the list of element on-disk and in-memory is the same

only_in_memory, only_on_disk = full_sdata._symmetric_difference_with_zarr_store()

assert len(only_in_memory) == 0

assert len(only_on_disk) == 0

full_sdata["new_image2d"] = full_sdata.images["image2d"]

full_sdata["new_labels2d"] = full_sdata.labels["labels2d"]

full_sdata["new_points_0"] = full_sdata.points["points_0"]

full_sdata["new_circles"] = full_sdata.shapes["circles"]

full_sdata["new_table"] = full_sdata.tables["table"]

del full_sdata.images["image2d"]

del full_sdata.labels["labels2d"]

del full_sdata.points["points_0"]

del full_sdata.shapes["circles"]

del full_sdata.tables["table"]

# now the list of element on-disk and in-memory is different

only_in_memory, only_on_disk = full_sdata._symmetric_difference_with_zarr_store()

assert set(only_in_memory) == {

"images/new_image2d",

"labels/new_labels2d",

"points/new_points_0",

"shapes/new_circles",

"tables/new_table",

}

assert set(only_on_disk) == {

"images/image2d",

"labels/labels2d",

"points/points_0",

"shapes/circles",

"tables/table",

}

@pytest.mark.filterwarnings("ignore:SpatialData is not stored in the most current format:UserWarning")

@pytest.mark.parametrize("sdata_container_format", SDATA_FORMATS)

def test_change_path_of_subset(full_sdata: SpatialData, sdata_container_format: SpatialDataContainerFormatType) -> None:

"""A subset SpatialData object has not Zarr path associated, show that we can reassign the path"""

# points_0_1 is the extra element that stays only on disk, satisfying the only_on_disk > 0 assertion.

full_sdata = full_sdata.subset(["image2d", "labels2d", "points_0", "circles", "table", "points_0_1"])

with tempfile.TemporaryDirectory() as tmpdir:

f = os.path.join(tmpdir, "data.zarr")

full_sdata.write(f, sdata_formats=sdata_container_format)

subset = full_sdata.subset(["image2d", "labels2d", "points_0", "circles", "table"])

assert subset.path is None

subset.path = Path(f)

assert subset.is_self_contained()

only_in_memory, only_on_disk = subset._symmetric_difference_with_zarr_store()

assert len(only_in_memory) == 0

assert len(only_on_disk) > 0

f2 = os.path.join(tmpdir, "data2.zarr")

subset.write(f2, sdata_formats=sdata_container_format)

assert subset.is_self_contained()

only_in_memory, only_on_disk = subset._symmetric_difference_with_zarr_store()

assert len(only_in_memory) == 0

assert len(only_on_disk) == 0

# changes in the subset object will be reflected in the original object (in-memory, on-disk only if we save

# them with .write_element())

scale = Scale([2.0], axes=("x",))

set_transformation(subset["image2d"], scale)

assert isinstance(get_transformation(full_sdata["image2d"]), Scale)

# if we don't want this, we can read the data again from the new path

sdata2 = SpatialData.read(f2)

set_transformation(sdata2["labels2d"], scale)

assert not isinstance(get_transformation(full_sdata["labels2d"]), Scale)

assert not isinstance(get_transformation(subset["labels2d"]), Scale)

def _check_valid_name(f: Callable[[str], Any]) -> None:

with pytest.raises(TypeError, match="Name must be a string, not "):

f(2)

with pytest.raises(ValueError, match="Name cannot be an empty string."):

f("")

with pytest.raises(ValueError, match="Name cannot be '.'"):

f(".")

with pytest.raises(ValueError, match="Name cannot be '..'"):