yaozarrs.write.v05#

def write_label(
    self,
    name: str,
    label_image: LabelImage,
    datasets: ArrayOrPyramid,
    *,
    progress: bool = False,
) -> Self:
    """Write a label immediately with its data.

    This method creates the label structure and writes data in one call.
    The labels group structure and LabelsGroup metadata are created/updated
    automatically. Use this for the "immediate write" workflow.

    Parameters
    ----------
    name : str
        Label name (becomes the subgroup path, e.g., "cells", "nuclei").
    label_image : LabelImage
        OME-Zarr LabelImage metadata model for this label.
    datasets : ArrayLike | Sequence[ArrayLike]
        Data array(s) for each resolution level. For a single dataset,
        pass the array directly without wrapping in a list.
    progress : bool, optional
        Show progress bar for dask arrays. Default is False.

    Returns
    -------
    Self
        The builder instance (for method chaining).

    Raises
    ------
    ValueError
        If a label with this name was already written or added.
    NotImplementedError
        If the LabelImage has multiple multiscales.
    """
    self._validate_label_name(name)

    # Initialize labels group structure if needed
    self._ensure_initialized()

    # Update labels/zarr.json with this label
    self._update_labels_group(name)

    # Write the label using the existing write_image function
    # (LabelImage is a subclass of Image)
    write_image(
        self._dest / name,
        label_image,
        datasets,
        writer=self._writer,
        chunks=self._chunks,
        shards=self._shards,
        overwrite=self._overwrite,
        compression=self._compression,
        progress=progress,
    )

    return self

def add_label(
    self,
    name: str,
    label_image: LabelImage,
    datasets: ShapeAndDTypeOrPyramid,
) -> Self:
    """Add a label for the prepare-only workflow.

    Registers a label to be created when `prepare()` is called. Use this
    when you want to create the Zarr structure without writing data
    immediately. After calling `prepare()`, write data to the returned
    array handles.

    Parameters
    ----------
    name : str
        Label name (becomes the subgroup path, e.g., "cells", "nuclei").
    label_image : LabelImage
        OME-Zarr LabelImage metadata model for this label.
    datasets : ShapeAndDTypeOrPyramid
        Shape/dtype spec(s) for each resolution level:
        - Single level: `(shape, dtype)`
        - Multiple levels: `[(shape1, dtype1), (shape2, dtype2)]`

    Returns
    -------
    Self
        The builder instance (for method chaining).

    Raises
    ------
    ValueError
        If a label with this name was already added or written, or if the
        number of specs doesn't match the metadata.
    NotImplementedError
        If the LabelImage has multiple multiscales.
    """
    self._validate_label_name(name)
    _, specs_seq = _validate_and_normalize_datasets(
        label_image, datasets, f"Label '{name}': "
    )
    self._labels[name] = (label_image, specs_seq)
    return self

def prepare(self) -> tuple[Path, dict[str, Any]]:
    """Create the Zarr hierarchy and return array handles.

    Creates the complete labels group structure including LabelsGroup
    metadata, and empty arrays for all registered labels. Call this after
    registering all labels with `add_label()`.

    The returned arrays support numpy-style indexing for writing data:
    `arrays["label_name/dataset"][:] = data`.

    Returns
    -------
    tuple[Path, dict[str, Any]]
        A tuple of (root_path, arrays) where `arrays` maps composite keys
        like `"cells/0"` (label name / dataset path) to array objects. The
        array type depends on the configured writer (zarr.Array or
        tensorstore.TensorStore).

    Raises
    ------
    ValueError
        If no labels have been added with `add_label()`.
    FileExistsError
        If destination exists and `overwrite` is False.
    ImportError
        If no suitable writer backend is installed.
    """
    if not self._labels:  # pragma: no cover
        raise ValueError("No labels added. Use add_label() before prepare().")

    # Merge existing labels with new labels
    all_labels = list(self._preexisting_labels)
    all_labels.extend([new for new in self._labels if new not in all_labels])

    # Create or update labels/zarr.json with LabelsGroup metadata
    # If labels group already exists, just update the metadata file with new labels
    labels_group = LabelsGroup(labels=all_labels)
    if self._dest.exists() and (self._dest / "zarr.json").exists():
        _update_zarr3_group(self._dest, labels_group)
    else:
        # Create new group
        _create_zarr3_group(self._dest, labels_group, self._overwrite)

    # Create arrays for each label using prepare_image
    all_arrays: dict[str, Any] = {}
    for label_name, (label_image, datasets) in self._labels.items():
        _label_path, label_arrays = prepare_image(
            self._dest / label_name,
            label_image,
            datasets,
            chunks=self._chunks,
            shards=self._shards,
            writer=self._writer,
            overwrite=self._overwrite,
            compression=self._compression,
        )
        # Flatten into all_arrays with "label_name/dataset" keys
        for dataset_path, arr in label_arrays.items():
            all_arrays[f"{label_name}/{dataset_path}"] = arr

    return self._dest, all_arrays

def write_well(
    self,
    row: str,
    col: str,
    images: Mapping[str, ImageWithDatasets],
    *,
    progress: bool = False,
) -> Self:
    """Write a well immediately with its `images` (fields of view) and data.

    This method creates the well structure and writes all field data in one
    call. The plate structure and well metadata are created/updated
    automatically. Plate metadata (rows, columns, wells) is auto-generated
    from all written wells and rewritten after each call.

    Parameters
    ----------
    row : str
        Row name like "A", "B", etc.
    col : str
        Column name like "1", "2", etc.
    images : Mapping[str, ImageWithDatasets]
        Mapping of `{fov -> (image_model, datasets)}` where:
        - fov: Field of view identifier like "0", "1", etc.
        - datasets can be:
          - Single array (for one dataset): `{"0": (image, data)}`
          - Sequence (for multiple datasets): `{"0": (image, [data1, data2])}`
    progress : bool, optional
        Show progress bar for dask arrays. Default is False.

    Returns
    -------
    Self
        The builder instance (for method chaining).

    Raises
    ------
    ValueError
        If row/col combination was already written or added, or if a user-
        provided Plate doesn't include this well.
    NotImplementedError
        If any Image has multiple multiscales.
    """
    # Validate well hasn't been used
    self._validate_well_coordinates(row, col)

    # Initialize plate structure if needed
    self._ensure_initialized()

    # Normalize fields (convert single arrays to sequences)
    normalized_fields: dict[str, tuple[Image, Sequence[ArrayLike]]] = {}
    for fov, (image_model, datasets) in images.items():
        _, datasets_seq = _validate_and_normalize_datasets(
            image_model, datasets, f"Well '{row}/{col}', field '{fov}': "
        )
        normalized_fields[fov] = (image_model, datasets_seq)

    # Track this well's data before writing
    self._written_wells_data[(row, col)] = cast(
        "dict[str, ImageWithDatasets]", normalized_fields
    )

    # Update plate metadata with the new well
    self._update_plate_metadata()

    # Generate Well metadata for this well and create well subgroup
    well_group_path = self._dest / f"{row}/{col}"
    well_metadata = self._generate_well_metadata(list(images))
    _create_zarr3_group(well_group_path, well_metadata, self._overwrite)

    # Write each field of view
    for fov, (image_model, datasets_seq) in normalized_fields.items():
        field_path = well_group_path / fov
        write_image(
            field_path,
            image_model,
            datasets_seq,
            writer=self._writer,
            chunks=self._chunks,
            shards=self._shards,
            overwrite=self._overwrite,
            compression=self._compression,
            progress=progress,
        )

    return self

def add_well(
    self,
    *,
    row: str,
    col: str,
    images: Mapping[str, ImageWithShapeSpecs],
) -> Self:
    """Add a well for the prepare-only workflow.

    Registers a well with its fields to be created when `prepare()` is called.
    Use this when you want to create the Zarr structure without writing data
    immediately. After calling `prepare()`, write data to the returned array
    handles.

    Parameters
    ----------
    row : str
        Row name like "A", "B", etc.
    col : str
        Column name like "1", "2", etc.
    images : Mapping[str, ImageWithShapeSpecs]
        Mapping of `{fov -> (image_model, specs)}` where specs provide the
        dtype and shape for each resolution level:
        - Single level: `(image, (shape, dtype))`
        - Multiple levels: `(image, [(shape1, dtype1), (shape2, dtype2)])`

    Returns
    -------
    Self
        The builder instance (for method chaining).

    Raises
    ------
    ValueError
        If row/col combination was already added/written, or if a user-
        provided Plate doesn't include this well, or if field specs
        don't match Image metadata.
    NotImplementedError
        If any Image has multiple multiscales.
    """
    # Validate well hasn't been used
    self._validate_well_coordinates(row, col)

    # Validate and normalize all fields before accepting
    well_path = f"{row}/{col}"
    normalized_fields: dict[str, ImageWithShapeSpecs] = {}

    for fov, (image_model, specs) in images.items():
        _, specs_seq = _validate_and_normalize_datasets(
            image_model, specs, f"Well '{well_path}', field '{fov}': "
        )
        normalized_fields[fov] = (image_model, specs_seq)

    self._wells[well_path] = normalized_fields
    return self

def prepare(self) -> tuple[Path, dict[str, Any]]:
    """Create the Zarr hierarchy and return array handles.

    Creates the complete Plate structure including plate metadata (auto-
    generated from registered wells), well subgroups with Well metadata,
    and empty arrays for all registered fields. Call this after registering
    all wells with `add_well()`.

    The returned arrays support numpy-style indexing for writing data:
    `arrays["well/field/dataset"][:] = data`.

    Returns
    -------
    tuple[Path, dict[str, Any]]
        A tuple of (root_path, arrays) where `arrays` maps composite keys
        like `"A/1/0/0"` (well_path / field / dataset_path) to array
        objects. The array type depends on the configured writer
        (zarr.Array or tensorstore.TensorStore).

    Raises
    ------
    ValueError
        If no wells have been added with `add_well()`.
    FileExistsError
        If destination exists and `overwrite` is False.
    ImportError
        If no suitable writer backend is installed.
    """
    if not self._wells:
        raise ValueError("No wells added. Use add_well() before prepare().")

    # Generate plate metadata from registered wells
    plate = _merge_plate_metadata(self._get_images_dict(), self._user_plate)

    # Create plate zarr.json
    _create_zarr3_group(
        self._dest,
        plate,
        self._overwrite,
        extra_attributes=self._extra_attributes,
    )

    # Create arrays for each well/field combination
    all_arrays: dict[str, Any] = {}

    for well_path, fields in self._wells.items():
        # Generate Well metadata and group
        well_metadata = self._generate_well_metadata(list(fields))
        well_group_path = self._dest / well_path
        _create_zarr3_group(well_group_path, well_metadata, self._overwrite)

        # Create arrays for each field
        for fov, (image_model, datasets) in fields.items():
            field_path = well_group_path / fov

            _field_path, field_arrays = prepare_image(
                field_path,
                image_model,
                datasets,
                chunks=self._chunks,
                shards=self._shards,
                writer=self._writer,
                overwrite=self._overwrite,
                compression=self._compression,
            )

            # Flatten into all_arrays with "well/field/dataset" keys
            for dataset_path, arr in field_arrays.items():
                composite_key = f"{well_path}/{fov}/{dataset_path}"
                all_arrays[composite_key] = arr

    return self._dest, all_arrays

def write_image(
    self,
    name: str,
    image: Image,
    datasets: ArrayOrPyramid,
    *,
    progress: bool = False,
) -> Self:
    """Write a series immediately with its data.

    This method creates the series structure and writes data in one call.
    The root structure and OME metadata are created/updated automatically.
    Use this for the "immediate write" workflow.

    Parameters
    ----------
    name : str
        Series name (becomes the subgroup path, e.g., "0", "1").
    image : Image
        OME-Zarr Image metadata model for this series.
    datasets : ArrayLike | Sequence[ArrayLike]
        Data array(s) for each resolution level. For a single dataset,
        pass the array directly without wrapping in a list.
    progress : bool, optional
        Show progress bar when writing dask arrays. Default is False.

    Returns
    -------
    Self
        The builder instance (for method chaining).

    Raises
    ------
    ValueError
        If a series with this name was already written or added.
    NotImplementedError
        If the Image has multiple multiscales.
    """
    self._validate_series_name(name)

    # Initialize root structure if needed
    self._ensure_initialized()

    # Update OME/zarr.json with this series
    self._update_ome_series(name)

    # Write the series using the existing write_image function
    write_image(
        self._dest / name,
        image,
        datasets,
        writer=self._writer,
        chunks=self._chunks,
        shards=self._shards,
        overwrite=self._overwrite,
        compression=self._compression,
        progress=progress,
    )

    return self

def add_series(
    self,
    name: str,
    image: Image,
    datasets: ShapeAndDTypeOrPyramid,
) -> Self:
    """Add a series for the prepare-only workflow.

    Registers a series to be created when `prepare()` is called. Use this
    when you want to create the Zarr structure without writing data
    immediately. After calling `prepare()`, write data to the returned
    array handles.

    Parameters
    ----------
    name : str
        Series name (becomes the subgroup path, e.g., "0", "1").
    image : Image
        OME-Zarr Image metadata model for this series.
    datasets : ShapeAndDType | Sequence[ShapeAndDType]
        Shape and dtype specification(s) for each resolution level, as
        `(shape, dtype)` tuples. For a single dataset, pass the tuple
        directly without wrapping in a list.

    Returns
    -------
    Self
        The builder instance (for method chaining).

    Raises
    ------
    ValueError
        If a series with this name was already added or written, or if the
        number of dataset specs doesn't match the metadata.
    NotImplementedError
        If the Image has multiple multiscales.
    """
    self._validate_series_name(name)
    _, datasets_seq = _validate_and_normalize_datasets(
        image, datasets, f"Series '{name}': "
    )
    self._series[name] = (image, datasets_seq)
    return self

def prepare(self) -> tuple[Path, dict[str, Any]]:
    """Create the Zarr hierarchy and return array handles.

    Creates the complete bioformats2raw structure including root metadata,
    OME directory with series list, and empty arrays for all registered
    series. Call this after registering all series with `add_series()`.

    The returned arrays support numpy-style indexing for writing data:
    `arrays["series/dataset"][:] = data`.

    Returns
    -------
    tuple[Path, dict[str, Any]]
        A tuple of (root_path, arrays) where `arrays` maps composite keys
        like `"0/0"` (series name / dataset path) to array objects. The
        array type depends on the configured writer (zarr.Array or
        tensorstore.TensorStore).

    Raises
    ------
    ValueError
        If no series have been added with `add_series()`.
    FileExistsError
        If destination exists and `overwrite` is False.
    ImportError
        If no suitable writer backend is installed.
    """
    if not self._series:  # pragma: no cover
        raise ValueError("No series added. Use add_series() before prepare().")

    # Create root zarr.json with bioformats2raw.layout
    bf2raw = Bf2Raw(bioformats2raw_layout=3)  # type: ignore
    _create_zarr3_group(
        self._dest,
        bf2raw,
        self._overwrite,
        extra_attributes=self._extra_attributes,
    )

    # Create OME/zarr.json with series list
    ome_path = self._dest / "OME"
    series_model = Series(series=list(self._series))
    _create_zarr3_group(ome_path, series_model, self._overwrite)

    # Write METADATA.ome.xml if provided
    if self._ome_xml is not None:
        (ome_path / "METADATA.ome.xml").write_text(self._ome_xml)

    # Create arrays for each series using prepare_image
    all_arrays: dict[str, Any] = {}
    for series_name, (image_model, dataset_specs) in self._series.items():
        _root_path, series_arrays = prepare_image(
            self._dest / series_name,
            image_model,
            dataset_specs,
            chunks=self._chunks,
            shards=self._shards,
            writer=self._writer,
            overwrite=self._overwrite,
            compression=self._compression,
        )
        # Flatten into all_arrays with "series/dataset" keys
        for dataset_path, arr in series_arrays.items():
            all_arrays[f"{series_name}/{dataset_path}"] = arr

    return self._dest, all_arrays