Refactor ManifestGroup (#518)

* improve type hinting by just using Mapping

* rename manifest array attribute

* add metadata property

* make ManifestGroup a Mapping

* rename manifest_group -> group

* typing

* improve the repr

* test new group API

Author: TomNicholas

virtualizarr/manifests/__init__.py

@@ -1,7 +1,7 @@
 # Note: This directory is named "manifests" rather than "manifest".
 # This is just to avoid conflicting with some type of file called manifest that .gitignore recommends ignoring.
 
-from .array import ManifestArray  # type: ignore # noqa
-from .group import ManifestGroup  # type: ignore # noqa
-from .manifest import ChunkEntry, ChunkManifest  # type: ignore # noqa
-from .store import ManifestStore  # type: ignore # noqa
+from virtualizarr.manifests.array import ManifestArray  # type: ignore # noqa
+from virtualizarr.manifests.group import ManifestGroup  # type: ignore # noqa
+from virtualizarr.manifests.manifest import ChunkEntry, ChunkManifest  # type: ignore # noqa
+from virtualizarr.manifests.store import ManifestStore  # type: ignore # noqa

virtualizarr/manifests/group.py

@@ -1,37 +1,107 @@
-from typing import Mapping, TypeAlias
+from __future__ import annotations
+
+import textwrap
+from typing import Iterator, Mapping
 
 from zarr.core.group import GroupMetadata
 
 from virtualizarr.manifests import ManifestArray
 
-ManifestArrayVariableMapping: TypeAlias = dict[str, ManifestArray]
-
 
-class ManifestGroup:
+class ManifestGroup(
+    Mapping[str, "ManifestArray | ManifestGroup"],
+):
     """
-    Virtualized representation of multiple ManifestArrays as a Zarr Group.
+    Immutable representation of a single virtual zarr group.
     """
 
-    # TODO: Consider refactoring according to https://github.com/zarr-developers/VirtualiZarr/pull/490#discussion_r2007805272
-    _manifest_arrays: Mapping[str, ManifestArray]
+    _members: Mapping[str, "ManifestArray | ManifestGroup"]
     _metadata: GroupMetadata
 
     def __init__(
         self,
-        manifest_arrays: ManifestArrayVariableMapping,
-        attributes: dict,
+        arrays: Mapping[str, ManifestArray] | None = None,
+        groups: Mapping[str, "ManifestGroup"] | None = None,
+        attributes: dict | None = None,
     ) -> None:
         """
-        Create a ManifestGroup from the dictionary of ManifestArrays and the group / dataset level metadata
+        Create a ManifestGroup containing ManifestArrays and/or sub-groups, as well as any group-level metadata.
 
         Parameters
         ----------
-        attributes : attributes to include in Group metadata
-        manifest_dict : ManifestArrayVariableMapping
+        arrays : Mapping[str, ManifestArray], optional
+            ManifestArray objects to represent virtual zarr arrays.
+        groups : Mapping[str, ManifestGroup], optional
+            ManifestGroup objects to represent virtual zarr subgroups.
+        attributes : dict, optional
+            Zarr attributes to add as zarr group metadata.
         """
-
         self._metadata = GroupMetadata(attributes=attributes)
-        self._manifest_arrays = manifest_arrays
 
-    def __str__(self) -> str:
-        return f"ManifestGroup(manifest_arrays={self._manifest_arrays}, metadata={self._metadata})"
+        _arrays: Mapping[str, ManifestArray] = {} if arrays is None else arrays
+
+        if groups:
+            # TODO add support for nested groups
+            raise NotImplementedError
+        else:
+            _groups: Mapping[str, ManifestGroup] = {} if groups is None else groups
+
+        for name, arr in _arrays.items():
+            if not isinstance(arr, ManifestArray):
+                raise TypeError(
+                    f"ManifestGroup can only wrap ManifestArray objects, but array {name} passed is of type {type(arr)}"
+                )
+
+        # TODO type check groups passed
+
+        # TODO check that all arrays have the same shapes or dimensions?
+        # Technically that's allowed by the zarr model, so we should theoretically only check that upon converting to xarray
+
+        colliding_names = set(_arrays.keys()).intersection(set(_groups.keys()))
+        if colliding_names:
+            raise ValueError(
+                f"Some names collide as they are present in both the array and group keys: {colliding_names}"
+            )
+
+        self._members = {**_arrays, **_groups}
+
+    @property
+    def metadata(self) -> GroupMetadata:
+        """Zarr group metadata."""
+        return self._metadata
+
+    @property
+    def arrays(self) -> dict[str, ManifestArray]:
+        """ManifestArrays contained in this group."""
+        return {k: v for k, v in self._members.items() if isinstance(v, ManifestArray)}
+
+    @property
+    def groups(self) -> dict[str, "ManifestGroup"]:
+        """Subgroups contained in this group."""
+        return {k: v for k, v in self._members.items() if isinstance(v, ManifestGroup)}
+
+    def __getitem__(self, path: str) -> "ManifestArray | ManifestGroup":
+        """Obtain a group member."""
+        if "/" in path:
+            raise ValueError(
+                f"ManifestGroup.__getitem__ can only be used to get immediate subgroups and subarrays, but received multi-part path {path}"
+            )
+
+        return self._members[path]
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._members.keys())
+
+    def __len__(self) -> int:
+        return len(self._members)
+
+    def __repr__(self) -> str:
+        return textwrap.dedent(
+            f"""
+            ManifestGroup(
+                arrays={self.arrays},
+                groups={self.groups},
+                metadata={self.metadata},
+            )
+            """
+        )

virtualizarr/manifests/store.py

@@ -2,7 +2,7 @@
 
 import pickle
 from collections.abc import Iterable
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Mapping
 from urllib.parse import urlparse
 
 from zarr.abc.store import (
@@ -15,6 +15,7 @@
 from zarr.core.buffer import Buffer
 from zarr.core.buffer.core import BufferPrototype
 
+from virtualizarr.manifests.array import ManifestArray
 from virtualizarr.manifests.group import ManifestGroup
 
 if TYPE_CHECKING:
@@ -27,6 +28,7 @@
 
 __all__ = ["ManifestStore"]
 
+
 _ALLOWED_EXCEPTIONS: tuple[type[Exception], ...] = (
     FileNotFoundError,
     IsADirectoryError,
@@ -39,7 +41,6 @@
 
 from zarr.core.buffer import default_buffer_prototype
 
-from virtualizarr.manifests.group import ManifestArrayVariableMapping
 from virtualizarr.vendor.zarr.metadata import dict_to_buffer
 
 if TYPE_CHECKING:
@@ -59,20 +60,20 @@ class StoreRequest:
 
 
 async def list_dir_from_manifest_arrays(
-    manifest_arrays: ManifestArrayVariableMapping, prefix: str
+    arrays: Mapping[str, ManifestArray], prefix: str
 ) -> AsyncGenerator[str]:
     """Create the expected results for Zarr's `store.list_dir()` from an Xarray DataArrray or Dataset
 
     Parameters
     ----------
-    manifest_arrays : ManifestArrayVariableMapping
+    arrays : Mapping[str, ManifestArrays]
     prefix : str
 
-
     Returns
     -------
     AsyncIterator[str]
     """
+    # TODO shouldn't this just accept a ManifestGroup instead?
     # Start with expected group level metadata
     raise NotImplementedError
 
@@ -99,11 +100,11 @@ def get_zarr_metadata(manifest_group: ManifestGroup, key: str) -> Buffer:
     # If requesting the root metadata, return the standard group metadata with additional dataset specific attributes
 
     if key == "zarr.json":
-        metadata = manifest_group._metadata.to_dict()
+        metadata = manifest_group.metadata.to_dict()
         return dict_to_buffer(metadata, prototype=default_buffer_prototype())
     else:
         var, _ = key.split("/")
-        metadata = manifest_group._manifest_arrays[var].metadata.to_dict()
+        metadata = manifest_group.arrays[var].metadata.to_dict()
         return dict_to_buffer(metadata, prototype=default_buffer_prototype())
 
 
@@ -170,16 +171,17 @@ def find_matching_store(stores: StoreDict, request_key: str) -> StoreRequest:
 
 
 class ManifestStore(Store):
-    """A read-only Zarr store that uses obstore to access data on AWS, GCP, Azure. The requests
+    """
+    A read-only Zarr store that uses obstore to access data on AWS, GCP, Azure. The requests
     from the Zarr API are redirected using the :class:`virtualizarr.manifests.ManifestGroup` containing
     multiple :class:`virtualizarr.manifests.ManifestArray`,
     allowing for virtually interfacing with underlying data in other file format.
 
-
     Parameters
     ----------
-    manifest_group : ManifestGroup
-        Manifest Group containing Group metadata and mapping variable names to ManifestArrays
+    group : ManifestGroup
+        Root group of the store.
+        Contains group metadata, ManifestArrays, and any subgroups.
     stores : dict[prefix, :class:`obstore.store.ObjectStore`]
         A mapping of url prefixes to obstore Store instances set up with the proper credentials.
 
@@ -196,15 +198,15 @@ class ManifestStore(Store):
     Modified from https://github.com/zarr-developers/zarr-python/pull/1661
     """
 
-    _manifest_group: ManifestGroup
+    _group: ManifestGroup
     _stores: StoreDict
 
     def __eq__(self, value: object):
         NotImplementedError
 
     def __init__(
         self,
-        manifest_group: ManifestGroup,
+        group: ManifestGroup,
         *,
         stores: StoreDict,  # TODO: Consider using a sequence of tuples rather than a dict (see https://github.com/zarr-developers/VirtualiZarr/pull/490#discussion_r2010717898).
     ) -> None:
@@ -223,14 +225,17 @@ def __init__(
         for store in stores.values():
             if not store.__class__.__module__.startswith("obstore"):
                 raise TypeError(f"expected ObjectStore class, got {store!r}")
+
         # TODO: Don't allow stores with prefix
-        # TODO: Type check the manifest arrays
+        if not isinstance(group, ManifestGroup):
+            raise TypeError
+
         super().__init__(read_only=True)
         self._stores = stores
-        self._manifest_group = manifest_group
+        self._group = group
 
     def __str__(self) -> str:
-        return f"ManifestStore({self._manifest_group}, {self._stores})"
+        return f"ManifestStore(group={self._group}, stores={self._stores})"
 
     def __getstate__(self) -> dict[Any, Any]:
         state = self.__dict__.copy()
@@ -257,10 +262,10 @@ async def get(
         import obstore as obs
 
         if key.endswith("zarr.json"):
-            return get_zarr_metadata(self._manifest_group, key)
+            return get_zarr_metadata(self._group, key)
         var, chunk_key = parse_manifest_index(key)
-        marr = self._manifest_group._manifest_arrays[var]
-        manifest = marr._manifest
+        marr = self._group.arrays[var]
+        manifest = marr.manifest
 
         path = manifest._paths[*chunk_key]
         offset = manifest._offsets[*chunk_key]
@@ -345,7 +350,7 @@ def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
     async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
         # docstring inherited
         yield "zarr.json"
-        for k in self._manifest_group._manifest_arrays.keys():
+        for k in self._group.arrays.keys():
             yield k
 
 

virtualizarr/tests/test_manifests/test_group.py

@@ -1,3 +1,5 @@
+import textwrap
+
 import pytest
 from zarr.core.group import GroupMetadata
 
@@ -20,17 +22,28 @@ def manifest_array(array_v3_metadata):
 
 
 class TestManifestGroup:
-    def test_manifest_array(self, array_v3_metadata, manifest_array):
+    def test_group_containing_array(self, manifest_array):
         var = "foo"
-        manifest_group = ManifestGroup(
-            manifest_arrays={var: manifest_array}, attributes={}
-        )
-        assert isinstance(manifest_group._manifest_arrays, dict)
-        assert isinstance(manifest_group._manifest_arrays[var], ManifestArray)
-        assert isinstance(manifest_group._metadata, GroupMetadata)
+        manifest_group = ManifestGroup(arrays={var: manifest_array}, attributes={})
+
+        assert manifest_group.arrays == {var: manifest_array}
+        assert manifest_group.groups == {}
+        assert isinstance(manifest_group[var], ManifestArray)
+        with pytest.raises(KeyError):
+            manifest_group["bar"]
+        assert isinstance(manifest_group.metadata, GroupMetadata)
+        assert len(manifest_group) == 1
+        assert list(manifest_group) == [var]
 
     def test_manifest_repr(self, manifest_array):
-        manifest_group = ManifestGroup(
-            manifest_arrays={"foo": manifest_array}, attributes={}
+        manifest_group = ManifestGroup(arrays={"foo": manifest_array}, attributes={})
+        expected_repr = textwrap.dedent(
+            """
+            ManifestGroup(
+                arrays={'foo': ManifestArray<shape=(5, 2, 20), dtype=int32, chunks=(5, 1, 10)>},
+                groups={},
+                metadata=GroupMetadata(attributes={}, zarr_format=3, consolidated_metadata=None, node_type='group'),
+            )
+            """
         )
-        assert str(manifest_group)
+        assert repr(manifest_group) == expected_repr

virtualizarr/tests/test_manifests/test_store.py

@@ -71,9 +71,10 @@ def _generate_manifest_store(
     array_metadata = array_v3_metadata(shape=shape, chunks=chunks)
     manifest_array = ManifestArray(metadata=array_metadata, chunkmanifest=manifest)
     manifest_group = ManifestGroup(
-        {"foo": manifest_array, "bar": manifest_array}, attributes={"Zarr": "Hooray!"}
+        arrays={"foo": manifest_array, "bar": manifest_array},
+        attributes={"Zarr": "Hooray!"},
     )
-    return ManifestStore(stores={prefix: store}, manifest_group=manifest_group)
+    return ManifestStore(stores={prefix: store}, group=manifest_group)
 
 
 @pytest.fixture()