fix:*-like creation routines take kwargs

changes/2992.fix.rst

@@ -0,0 +1,3 @@
+Fix a bug preventing ``ones_like``, ``full_like``, ``empty_like``, ``zeros_like`` and ``open_like`` functions from accepting
+an explicit specification of array attributes like shape, dtype, chunks etc. The functions ``full_like``,
+``empty_like``, and ``open_like`` now also more consistently infer a ``fill_value`` parameter from the provided array.

src/zarr/api/asynchronous.py

@@ -3,7 +3,7 @@
 import asyncio
 import dataclasses
 import warnings
-from typing import TYPE_CHECKING, Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal, NotRequired, TypedDict, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -47,9 +47,12 @@
 if TYPE_CHECKING:
     from collections.abc import Iterable
 
+    import numcodecs
+
     from zarr.abc.codec import Codec
     from zarr.core.buffer import NDArrayLikeOrScalar
     from zarr.core.chunk_key_encodings import ChunkKeyEncoding
+    from zarr.core.metadata.v2 import CompressorLikev2
     from zarr.storage import StoreLike
 
     # TODO: this type could use some more thought
@@ -118,10 +121,20 @@
     return shape, chunks
 
 
-def _like_args(a: ArrayLike, kwargs: dict[str, Any]) -> dict[str, Any]:
+class _LikeArgs(TypedDict):
+    shape: NotRequired[ChunkCoords]
+    chunks: NotRequired[ChunkCoords]
+    dtype: NotRequired[np.dtype[np.generic]]
+    order: NotRequired[Literal["C", "F"]]
+    filters: NotRequired[tuple[numcodecs.abc.Codec, ...] | None]
+    compressor: NotRequired[CompressorLikev2]
+    codecs: NotRequired[tuple[Codec, ...]]
+
+
+def _like_args(a: ArrayLike) -> _LikeArgs:
     """Set default values for shape and chunks if they are not present in the array-like object"""
 
-    new = kwargs.copy()
+    new: _LikeArgs = {}
 
     shape, chunks = _get_shape_chunks(a)
     if shape is not None:
@@ -1088,7 +1101,7 @@
     shape: ChunkCoords, **kwargs: Any
 ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an empty array with the specified shape. The contents will be filled with the
-    array's fill value or zeros if no fill value is provided.
+    specified fill value or zeros if no fill value is provided.
 
     Parameters
     ----------
@@ -1103,8 +1116,7 @@
     retrieve data from an empty Zarr array, any values may be returned,
     and these are not guaranteed to be stable from one access to the next.
     """
-
-    return await create(shape=shape, fill_value=None, **kwargs)
+    return await create(shape=shape, **kwargs)
 
 
 async def empty_like(
@@ -1131,8 +1143,10 @@
     retrieve data from an empty Zarr array, any values may be returned,
     and these are not guaranteed to be stable from one access to the next.
     """
-    like_kwargs = _like_args(a, kwargs)
-    return await empty(**like_kwargs)
+    like_kwargs = _like_args(a) | kwargs
+    if isinstance(a, (AsyncArray | Array)):
+        like_kwargs.setdefault("fill_value", a.metadata.fill_value)
+    return await empty(**like_kwargs)  # type: ignore[arg-type]
 
 
 # TODO: add type annotations for fill_value and kwargs
@@ -1177,10 +1191,10 @@
     Array
         The new array.
     """
-    like_kwargs = _like_args(a, kwargs)
-    if isinstance(a, AsyncArray):
+    like_kwargs = _like_args(a) | kwargs
+    if isinstance(a, (AsyncArray | Array)):
         like_kwargs.setdefault("fill_value", a.metadata.fill_value)
-    return await full(**like_kwargs)
+    return await full(**like_kwargs)  # type: ignore[arg-type]
 
 
 async def ones(
@@ -1221,8 +1235,8 @@
     Array
         The new array.
     """
-    like_kwargs = _like_args(a, kwargs)
-    return await ones(**like_kwargs)
+    like_kwargs = _like_args(a) | kwargs
+    return await ones(**like_kwargs)  # type: ignore[arg-type]
 
 
 async def open_array(
@@ -1302,10 +1316,10 @@
     AsyncArray
         The opened array.
     """
-    like_kwargs = _like_args(a, kwargs)
+    like_kwargs = _like_args(a) | kwargs
     if isinstance(a, (AsyncArray | Array)):
-        kwargs.setdefault("fill_value", a.metadata.fill_value)
-    return await open_array(path=path, **like_kwargs)
+        like_kwargs.setdefault("fill_value", a.metadata.fill_value)
+    return await open_array(path=path, **like_kwargs)  # type: ignore[arg-type]
 
 
 async def zeros(
@@ -1346,5 +1360,5 @@
     Array
         The new array.
     """
-    like_kwargs = _like_args(a, kwargs)
-    return await zeros(**like_kwargs)
+    like_kwargs = _like_args(a) | kwargs
+    return await zeros(**like_kwargs)  # type: ignore[arg-type]

tests/test_api.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import re
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import zarr.codecs
 import zarr.storage
@@ -77,6 +77,91 @@ def test_create(memory_store: Store) -> None:
         z = create(shape=(400, 100), chunks=(16, 16.5), store=store, overwrite=True)  # type: ignore [arg-type]
 
 
+@pytest.mark.parametrize(
+    "func",
+    [
+        zarr.api.asynchronous.zeros_like,
+        zarr.api.asynchronous.ones_like,
+        zarr.api.asynchronous.empty_like,
+        zarr.api.asynchronous.full_like,
+        zarr.api.asynchronous.open_like,
+    ],
+)
+@pytest.mark.parametrize("out_shape", ["keep", (10, 10)])
+@pytest.mark.parametrize("out_chunks", ["keep", (10, 10)])
+@pytest.mark.parametrize("out_dtype", ["keep", "int8"])
+@pytest.mark.parametrize("out_fill", ["keep", 4])
+async def test_array_like_creation(
+    zarr_format: ZarrFormat,
+    func: Callable[[Any], Any],
+    out_shape: Literal["keep"] | tuple[int, ...],
+    out_chunks: Literal["keep"] | tuple[int, ...],
+    out_dtype: str,
+    out_fill: Literal["keep"] | int,
+) -> None:
+    """
+    Test zeros_like, ones_like, empty_like, full_like, ensuring that we can override the
+    shape, chunks, dtype and fill_value of the array-like object provided to these functions with
+    appropriate keyword arguments
+    """
+    ref_fill = 100
+    ref_arr = zarr.create_array(
+        store={},
+        shape=(11, 12),
+        dtype="uint8",
+        chunks=(11, 12),
+        zarr_format=zarr_format,
+        fill_value=ref_fill,
+    )
+    kwargs: dict[str, object] = {}
+    if func is zarr.api.asynchronous.full_like:
+        if out_fill == "keep":
+            expect_fill = ref_fill
+        else:
+            expect_fill = out_fill
+            kwargs["fill_value"] = expect_fill
+    elif func is zarr.api.asynchronous.zeros_like:
+        expect_fill = 0
+    elif func is zarr.api.asynchronous.ones_like:
+        expect_fill = 1
+    elif func is zarr.api.asynchronous.empty_like:
+        if out_fill == "keep":
+            expect_fill = ref_fill
+        else:
+            kwargs["fill_value"] = out_fill
+            expect_fill = out_fill
+    elif func is zarr.api.asynchronous.open_like:  # type: ignore[comparison-overlap]
+        if out_fill == "keep":
+            expect_fill = ref_fill
+        else:
+            kwargs["fill_value"] = out_fill
+            expect_fill = out_fill
+        kwargs["mode"] = "w"
+    else:
+        raise AssertionError
+    if out_shape != "keep":
+        kwargs["shape"] = out_shape
+        expect_shape = out_shape
+    else:
+        expect_shape = ref_arr.shape
+    if out_chunks != "keep":
+        kwargs["chunks"] = out_chunks
+        expect_chunks = out_chunks
+    else:
+        expect_chunks = ref_arr.chunks
+    if out_dtype != "keep":
+        kwargs["dtype"] = out_dtype
+        expect_dtype = out_dtype
+    else:
+        expect_dtype = ref_arr.dtype  # type: ignore[assignment]
+
+    new_arr = await func(ref_arr, path="foo", **kwargs)  # type: ignore[call-arg]
+    assert new_arr.shape == expect_shape
+    assert new_arr.chunks == expect_chunks
+    assert new_arr.dtype == expect_dtype
+    assert np.all(Array(new_arr)[:] == expect_fill)
+
+
 # TODO: parametrize over everything this function takes
 @pytest.mark.parametrize("store", ["memory"], indirect=True)
 def test_create_array(store: Store, zarr_format: ZarrFormat) -> None:

tests/test_group.py

@@ -7,7 +7,7 @@
 import re
 import time
 import warnings
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, get_args
 
 import numpy as np
 import pytest
@@ -670,6 +670,66 @@ def test_group_create_array(
     assert np.array_equal(array[:], data)
 
 
+LikeMethodName = Literal["zeros_like", "ones_like", "empty_like", "full_like"]
+
+
+@pytest.mark.parametrize("method_name", get_args(LikeMethodName))
+@pytest.mark.parametrize("out_shape", ["keep", (10, 10)])
+@pytest.mark.parametrize("out_chunks", ["keep", (10, 10)])
+@pytest.mark.parametrize("out_dtype", ["keep", "int8"])
+def test_group_array_like_creation(
+    zarr_format: ZarrFormat,
+    method_name: LikeMethodName,
+    out_shape: Literal["keep"] | tuple[int, ...],
+    out_chunks: Literal["keep"] | tuple[int, ...],
+    out_dtype: str,
+) -> None:
+    """
+    Test Group.{zeros_like, ones_like, empty_like, full_like}, ensuring that we can override the
+    shape, chunks, and dtype of the array-like object provided to these functions with
+    appropriate keyword arguments
+    """
+    ref_arr = zarr.ones(store={}, shape=(11, 12), dtype="uint8", chunks=(11, 12))
+    group = Group.from_store({}, zarr_format=zarr_format)
+    kwargs = {}
+    if method_name == "full_like":
+        expect_fill = 4
+        kwargs["fill_value"] = expect_fill
+        meth = group.full_like
+    elif method_name == "zeros_like":
+        expect_fill = 0
+        meth = group.zeros_like
+    elif method_name == "ones_like":
+        expect_fill = 1
+        meth = group.ones_like
+    elif method_name == "empty_like":
+        expect_fill = ref_arr.fill_value
+        meth = group.empty_like
+    else:
+        raise AssertionError
+    if out_shape != "keep":
+        kwargs["shape"] = out_shape
+        expect_shape = out_shape
+    else:
+        expect_shape = ref_arr.shape
+    if out_chunks != "keep":
+        kwargs["chunks"] = out_chunks
+        expect_chunks = out_chunks
+    else:
+        expect_chunks = ref_arr.chunks
+    if out_dtype != "keep":
+        kwargs["dtype"] = out_dtype
+        expect_dtype = out_dtype
+    else:
+        expect_dtype = ref_arr.dtype
+
+    new_arr = meth(name="foo", data=ref_arr, **kwargs)
+    assert new_arr.shape == expect_shape
+    assert new_arr.chunks == expect_chunks
+    assert new_arr.dtype == expect_dtype
+    assert np.all(new_arr[:] == expect_fill)
+
+
 def test_group_array_creation(
     store: Store,
     zarr_format: ZarrFormat,