Skip to content

Commit

Permalink
Feature: group and array name properties (#1940)
Browse files Browse the repository at this point in the history
* feature: group and array path/name/basename properties

* tests
  • Loading branch information
jhamman authored Jun 3, 2024
1 parent 24e855c commit 72005d7
Show file tree
Hide file tree
Showing 4 changed files with 131 additions and 0 deletions.
38 changes: 38 additions & 0 deletions src/zarr/array.py
Original file line number Diff line number Diff line change
Expand Up @@ -378,6 +378,29 @@ def dtype(self) -> np.dtype[Any]:
def attrs(self) -> dict[str, JSON]:
return self.metadata.attributes

@property
def path(self) -> str:
"""Storage path."""
return self.store_path.path

@property
def name(self) -> str | None:
"""Array name following h5py convention."""
if self.path:
# follow h5py convention: add leading slash
name = self.path
if name[0] != "/":
name = "/" + name
return name
return None

@property
def basename(self) -> str | None:
"""Final component of name."""
if self.name is not None:
return self.name.split("/")[-1]
return None

async def _get_selection(
self,
indexer: Indexer,
Expand Down Expand Up @@ -630,6 +653,21 @@ def dtype(self) -> np.dtype[Any]:
def attrs(self) -> Attributes:
return Attributes(self)

@property
def path(self) -> str:
"""Storage path."""
return self._async_array.path

@property
def name(self) -> str | None:
"""Array name following h5py convention."""
return self._async_array.name

@property
def basename(self) -> str | None:
"""Final component of name."""
return self._async_array.basename

@property
def metadata(self) -> ArrayMetadata:
return self._async_array.metadata
Expand Down
38 changes: 38 additions & 0 deletions src/zarr/group.py
Original file line number Diff line number Diff line change
Expand Up @@ -270,6 +270,27 @@ async def _save_metadata(self) -> None:
awaitables = [set_or_delete(self.store_path / key, value) for key, value in to_save.items()]
await asyncio.gather(*awaitables)

@property
def path(self) -> str:
"""Storage path."""
return self.store_path.path

@property
def name(self) -> str:
"""Group name following h5py convention."""
if self.path:
# follow h5py convention: add leading slash
name = self.path
if name[0] != "/":
name = "/" + name
return name
return "/"

@property
def basename(self) -> str:
"""Final component of name."""
return self.name.split("/")[-1]

@property
def attrs(self) -> dict[str, Any]:
return self.metadata.attributes
Expand Down Expand Up @@ -462,13 +483,15 @@ def create(
store: StoreLike,
*,
attributes: dict[str, Any] = {}, # noqa: B006, FIXME
zarr_format: ZarrFormat = 3,
exists_ok: bool = False,
) -> Group:
obj = sync(
AsyncGroup.create(
store,
attributes=attributes,
exists_ok=exists_ok,
zarr_format=zarr_format,
),
)

Expand Down Expand Up @@ -521,6 +544,21 @@ def store_path(self) -> StorePath:
def metadata(self) -> GroupMetadata:
return self._async_group.metadata

@property
def path(self) -> str:
"""Storage path."""
return self._async_group.path

@property
def name(self) -> str:
"""Group name following h5py convention."""
return self._async_group.name

@property
def basename(self) -> str:
"""Final component of name."""
return self._async_group.basename

@property
def attrs(self) -> Attributes:
return Attributes(self)
Expand Down
36 changes: 36 additions & 0 deletions tests/v3/test_array.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
import pytest

from zarr.array import Array
from zarr.common import ZarrFormat
from zarr.group import Group
from zarr.store import LocalStore, MemoryStore


@pytest.mark.parametrize("store", ("local", "memory"), indirect=["store"])
@pytest.mark.parametrize("zarr_format", (2, 3))
def test_array_name_properties_no_group(
store: LocalStore | MemoryStore, zarr_format: ZarrFormat
) -> None:
arr = Array.create(store=store, shape=(100,), chunks=(10,), zarr_format=zarr_format, dtype="i4")
assert arr.path == ""
assert arr.name is None
assert arr.basename is None


@pytest.mark.parametrize("store", ("local", "memory"), indirect=["store"])
@pytest.mark.parametrize("zarr_format", (2, 3))
def test_array_name_properties_with_group(
store: LocalStore | MemoryStore, zarr_format: ZarrFormat
) -> None:
root = Group.create(store=store, zarr_format=zarr_format)
foo = root.create_array("foo", shape=(100,), chunks=(10,), dtype="i4")
assert foo.path == "foo"
assert foo.name == "/foo"
assert foo.basename == "foo"

bar = root.create_group("bar")
spam = bar.create_array("spam", shape=(100,), chunks=(10,), dtype="i4")

assert spam.path == "bar/spam"
assert spam.name == "/bar/spam"
assert spam.basename == "spam"
19 changes: 19 additions & 0 deletions tests/v3/test_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -372,3 +372,22 @@ def test_group_init(store: LocalStore | MemoryStore, zarr_format: ZarrFormat) ->
agroup = sync(AsyncGroup.create(store=store, zarr_format=zarr_format))
group = Group(agroup)
assert group._async_group == agroup


@pytest.mark.parametrize("store", ("local", "memory"), indirect=["store"])
@pytest.mark.parametrize("zarr_format", (2, 3))
def test_group_name_properties(store: LocalStore | MemoryStore, zarr_format: ZarrFormat) -> None:
root = Group.create(store=store, zarr_format=zarr_format)
assert root.path == ""
assert root.name == "/"
assert root.basename == ""

foo = root.create_group("foo")
assert foo.path == "foo"
assert foo.name == "/foo"
assert foo.basename == "foo"

bar = root.create_group("foo/bar")
assert bar.path == "foo/bar"
assert bar.name == "/foo/bar"
assert bar.basename == "bar"

0 comments on commit 72005d7

Please sign in to comment.