Repository: dask/dask
File: dask/array/core.py (Python)

Code before:
        )
        # We almost certainly get here because auto chunking has been used
        # on irregular chunks. The max will then be smaller than auto, so using
        # max is a safe choice
        arr = arr.rechunk(tuple(map(max, arr.chunks)))

    if region is not None:
        raise ValueError("Cannot use `region` keyword when url is not a `zarr.Array`.")

    zarr_read_kwargs = {} if zarr_read_kwargs is None else dict(zarr_read_kwargs)
    zarr_store = _setup_zarr_store(url, storage_options, **zarr_read_kwargs)

    zarr_array_kwargs.setdefault("shape", arr.shape)
    zarr_array_kwargs.setdefault("chunks", tuple(c[0] for c in arr.chunks))
    zarr_array_kwargs.setdefault("dtype", arr.dtype)

    array_name = component or zarr_array_kwargs.pop("name", None)
    if _zarr_v3():
        root = zarr.open_group(store=zarr_store, mode="a") if array_name else None
        if array_name:
            z = root.create_array(name=array_name, **zarr_array_kwargs)
        else:
            zarr_array_kwargs["store"] = zarr_store
            z = zarr.create_array(**zarr_array_kwargs)
    else:
        # TODO: drop this as soon as zarr v2 gets dropped.
        z = zarr.create(
            store=zarr_store,
            path=array_name,
            **zarr_array_kwargs,
        )

    return arr.store(z, lock=False, compute=compute, return_stored=return_stored)


def _get_zarr_write_chunks(zarr_array) -> tuple[int, ...]:
    """Get the appropriate chunk shape for writing to a Zarr array.

    For Zarr v3 arrays with sharding, returns the shard shape.
    For arrays without sharding, returns the chunk shape.
    For Zarr v2 arrays, returns the chunk shape.

    Parameters
    ----------
    zarr_array : zarr.Array
        The target zarr array

    Returns
    -------
    tuple
        The chunk shape to use for rechunking the dask array

Review comment (suggestion):
```suggestion
        # TODO: drop this as soon as zarr v2 gets dropped.
        # https://github.com/dask/dask/issues/12188
```

Code after:
        )
        # We almost certainly get here because auto chunking has been used
        # on irregular chunks. The max will then be smaller than auto, so using
        # max is a safe choice
        arr = arr.rechunk(tuple(map(max, arr.chunks)))

    if region is not None:
        raise ValueError("Cannot use `region` keyword when url is not a `zarr.Array`.")

    zarr_read_kwargs = {} if zarr_read_kwargs is None else dict(zarr_read_kwargs)
    zarr_store = _setup_zarr_store(url, storage_options, **zarr_read_kwargs)

    zarr_array_kwargs.setdefault("shape", arr.shape)
    zarr_array_kwargs.setdefault("chunks", tuple(c[0] for c in arr.chunks))
    zarr_array_kwargs.setdefault("dtype", arr.dtype)

    array_name = component or zarr_array_kwargs.pop("name", None)
    if _zarr_v3():
        root = zarr.open_group(store=zarr_store, mode="a") if array_name else None
        if array_name:
            z = root.create_array(name=array_name, **zarr_array_kwargs)
        else:
            zarr_array_kwargs["store"] = zarr_store
            z = zarr.create_array(**zarr_array_kwargs)
    else:
        # TODO: drop this as soon as zarr v2 gets dropped.
        # https://github.com/dask/dask/issues/12188
        z = zarr.create(
            store=zarr_store,
            path=array_name,
            **zarr_array_kwargs,
        )

    return arr.store(z, lock=False, compute=compute, return_stored=return_stored)


def _get_zarr_write_chunks(zarr_array) -> tuple[int, ...]:
    """Get the appropriate chunk shape for writing to a Zarr array.

    For Zarr v3 arrays with sharding, returns the shard shape.
    For arrays without sharding, returns the chunk shape.
    For Zarr v2 arrays, returns the chunk shape.

    Parameters
    ----------
    zarr_array : zarr.Array
        The target zarr array

    Returns
    -------
    tuple