Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[v3] Implement Group methods for empty, full, ones, and zeros #2210

Merged
merged 11 commits into from
Sep 19, 2024
4 changes: 2 additions & 2 deletions src/zarr/api/asynchronous.py
Original file line number Diff line number Diff line change
Expand Up @@ -575,7 +575,7 @@ async def create(
chunks: ChunkCoords | None = None, # TODO: v2 allowed chunks=True
dtype: npt.DTypeLike | None = None,
compressor: dict[str, JSON] | None = None, # TODO: default and type change
fill_value: Any = 0, # TODO: need type
fill_value: Any | None = 0, # TODO: need type
order: MemoryOrder | None = None, # TODO: default change
store: str | StoreLike | None = None,
synchronizer: Any | None = None,
Expand Down Expand Up @@ -827,7 +827,7 @@ async def full_like(a: ArrayLike, **kwargs: Any) -> AsyncArray:
"""
like_kwargs = _like_args(a, kwargs)
if isinstance(a, AsyncArray):
kwargs.setdefault("fill_value", a.metadata.fill_value)
like_kwargs.setdefault("fill_value", a.metadata.fill_value)
return await full(**like_kwargs)


Expand Down
91 changes: 59 additions & 32 deletions src/zarr/core/group.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import numpy.typing as npt
from typing_extensions import deprecated

import zarr.api.asynchronous as async_api
from zarr.abc.metadata import Metadata
from zarr.abc.store import set_or_delete
from zarr.core.array import Array, AsyncArray
Expand Down Expand Up @@ -704,29 +705,41 @@ async def arrays(self) -> AsyncGenerator[AsyncArray, None]:
async def tree(self, expand: bool = False, level: int | None = None) -> Any:
raise NotImplementedError

async def empty(self, **kwargs: Any) -> AsyncArray:
raise NotImplementedError
async def empty(self, *, name: str, shape: ChunkCoords, **kwargs: Any) -> AsyncArray:
return await async_api.empty(shape=shape, store=self.store_path, path=name, **kwargs)

async def zeros(self, **kwargs: Any) -> AsyncArray:
raise NotImplementedError
async def zeros(self, *, name: str, shape: ChunkCoords, **kwargs: Any) -> AsyncArray:
return await async_api.zeros(shape=shape, store=self.store_path, path=name, **kwargs)

async def ones(self, **kwargs: Any) -> AsyncArray:
raise NotImplementedError
async def ones(self, *, name: str, shape: ChunkCoords, **kwargs: Any) -> AsyncArray:
return await async_api.ones(shape=shape, store=self.store_path, path=name, **kwargs)

async def full(self, **kwargs: Any) -> AsyncArray:
raise NotImplementedError
async def full(
self, *, name: str, shape: ChunkCoords, fill_value: Any | None, **kwargs: Any
) -> AsyncArray:
return await async_api.full(
shape=shape, fill_value=fill_value, store=self.store_path, path=name, **kwargs
)

async def empty_like(self, prototype: AsyncArray, **kwargs: Any) -> AsyncArray:
raise NotImplementedError
async def empty_like(
self, *, name: str, prototype: async_api.ArrayLike, **kwargs: Any
) -> AsyncArray:
return await async_api.empty_like(a=prototype, store=self.store_path, path=name, **kwargs)

async def zeros_like(self, prototype: AsyncArray, **kwargs: Any) -> AsyncArray:
raise NotImplementedError
async def zeros_like(
self, *, name: str, prototype: async_api.ArrayLike, **kwargs: Any
) -> AsyncArray:
return await async_api.zeros_like(a=prototype, store=self.store_path, path=name, **kwargs)

async def ones_like(self, prototype: AsyncArray, **kwargs: Any) -> AsyncArray:
raise NotImplementedError
async def ones_like(
self, *, name: str, prototype: async_api.ArrayLike, **kwargs: Any
) -> AsyncArray:
return await async_api.ones_like(a=prototype, store=self.store_path, path=name, **kwargs)

async def full_like(self, prototype: AsyncArray, **kwargs: Any) -> AsyncArray:
raise NotImplementedError
async def full_like(
self, *, name: str, prototype: async_api.ArrayLike, **kwargs: Any
) -> AsyncArray:
return await async_api.full_like(a=prototype, store=self.store_path, path=name, **kwargs)

async def move(self, source: str, dest: str) -> None:
raise NotImplementedError
Expand Down Expand Up @@ -1058,29 +1071,43 @@ def require_array(self, name: str, **kwargs: Any) -> Array:
"""
return Array(self._sync(self._async_group.require_array(name, **kwargs)))

def empty(self, **kwargs: Any) -> Array:
return Array(self._sync(self._async_group.empty(**kwargs)))
def empty(self, *, name: str, shape: ChunkCoords, **kwargs: Any) -> Array:
return Array(self._sync(self._async_group.empty(name=name, shape=shape, **kwargs)))

def zeros(self, **kwargs: Any) -> Array:
return Array(self._sync(self._async_group.zeros(**kwargs)))
def zeros(self, *, name: str, shape: ChunkCoords, **kwargs: Any) -> Array:
return Array(self._sync(self._async_group.zeros(name=name, shape=shape, **kwargs)))

def ones(self, **kwargs: Any) -> Array:
return Array(self._sync(self._async_group.ones(**kwargs)))
def ones(self, *, name: str, shape: ChunkCoords, **kwargs: Any) -> Array:
return Array(self._sync(self._async_group.ones(name=name, shape=shape, **kwargs)))

def full(self, **kwargs: Any) -> Array:
return Array(self._sync(self._async_group.full(**kwargs)))
def full(
self, *, name: str, shape: ChunkCoords, fill_value: Any | None, **kwargs: Any
) -> Array:
return Array(
self._sync(
self._async_group.full(name=name, shape=shape, fill_value=fill_value, **kwargs)
)
)

def empty_like(self, prototype: AsyncArray, **kwargs: Any) -> Array:
return Array(self._sync(self._async_group.empty_like(prototype, **kwargs)))
def empty_like(self, *, name: str, prototype: async_api.ArrayLike, **kwargs: Any) -> Array:
return Array(
self._sync(self._async_group.empty_like(name=name, prototype=prototype, **kwargs))
)

def zeros_like(self, prototype: AsyncArray, **kwargs: Any) -> Array:
return Array(self._sync(self._async_group.zeros_like(prototype, **kwargs)))
def zeros_like(self, *, name: str, prototype: async_api.ArrayLike, **kwargs: Any) -> Array:
return Array(
self._sync(self._async_group.zeros_like(name=name, prototype=prototype, **kwargs))
)

def ones_like(self, prototype: AsyncArray, **kwargs: Any) -> Array:
return Array(self._sync(self._async_group.ones_like(prototype, **kwargs)))
def ones_like(self, *, name: str, prototype: async_api.ArrayLike, **kwargs: Any) -> Array:
return Array(
self._sync(self._async_group.ones_like(name=name, prototype=prototype, **kwargs))
)

def full_like(self, prototype: AsyncArray, **kwargs: Any) -> Array:
return Array(self._sync(self._async_group.full_like(prototype, **kwargs)))
def full_like(self, *, name: str, prototype: async_api.ArrayLike, **kwargs: Any) -> Array:
return Array(
self._sync(self._async_group.full_like(name=name, prototype=prototype, **kwargs))
)

def move(self, source: str, dest: str) -> None:
return self._sync(self._async_group.move(source, dest))
Expand Down
67 changes: 67 additions & 0 deletions tests/v3/test_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -390,6 +390,73 @@ def test_group_create_array(
assert np.array_equal(array[:], data)


def test_group_array_creation(
store: Store,
zarr_format: ZarrFormat,
):
group = Group.create(store, zarr_format=zarr_format)
shape = (10, 10)
empty_array = group.empty(name="empty", shape=shape)
assert isinstance(empty_array, Array)
assert empty_array.fill_value == 0
assert empty_array.shape == shape
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

assert empty_array.store_path.store == store

empty_like_array = group.empty_like(name="empty_like", prototype=empty_array)
assert isinstance(empty_like_array, Array)
assert empty_like_array.fill_value == 0
assert empty_like_array.shape == shape
jhamman marked this conversation as resolved.
Show resolved Hide resolved
assert empty_like_array.store_path.store == store

empty_array_bool = group.empty(name="empty_bool", shape=shape, dtype=np.dtype("bool"))
assert isinstance(empty_array_bool, Array)
assert not empty_array_bool.fill_value
assert empty_array_bool.shape == shape
assert empty_array_bool.store_path.store == store

empty_like_array_bool = group.empty_like(name="empty_like_bool", prototype=empty_array_bool)
assert isinstance(empty_like_array_bool, Array)
assert not empty_like_array_bool.fill_value
assert empty_like_array_bool.shape == shape
assert empty_like_array_bool.store_path.store == store

zeros_array = group.zeros(name="zeros", shape=shape)
assert isinstance(zeros_array, Array)
assert zeros_array.fill_value == 0
assert zeros_array.shape == shape
assert zeros_array.store_path.store == store

zeros_like_array = group.zeros_like(name="zeros_like", prototype=zeros_array)
assert isinstance(zeros_like_array, Array)
assert zeros_like_array.fill_value == 0
assert zeros_like_array.shape == shape
assert zeros_like_array.store_path.store == store

ones_array = group.ones(name="ones", shape=shape)
assert isinstance(ones_array, Array)
assert ones_array.fill_value == 1
assert ones_array.shape == shape
assert ones_array.store_path.store == store

ones_like_array = group.ones_like(name="ones_like", prototype=ones_array)
assert isinstance(ones_like_array, Array)
assert ones_like_array.fill_value == 1
assert ones_like_array.shape == shape
assert ones_like_array.store_path.store == store

full_array = group.full(name="full", shape=shape, fill_value=42)
assert isinstance(full_array, Array)
assert full_array.fill_value == 42
assert full_array.shape == shape
assert full_array.store_path.store == store

full_like_array = group.full_like(name="full_like", prototype=full_array, fill_value=43)
assert isinstance(full_like_array, Array)
assert full_like_array.fill_value == 43
assert full_like_array.shape == shape
assert full_like_array.store_path.store == store


@pytest.mark.parametrize("store", ("local", "memory", "zip"), indirect=["store"])
@pytest.mark.parametrize("zarr_format", (2, 3))
@pytest.mark.parametrize("exists_ok", [True, False])
Expand Down