Skip to content

Commit

Permalink
fix: define utility for converting asyncarray to array, and similar f…
Browse files Browse the repository at this point in the history
…or group, largely to appease mypy
  • Loading branch information
d-v-b committed May 10, 2024
1 parent 8e61ff9 commit b996aff
Showing 1 changed file with 25 additions and 5 deletions.
30 changes: 25 additions & 5 deletions src/zarr/group.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
from zarr.config import RuntimeConfiguration, SyncConfiguration
from zarr.store import StoreLike, StorePath, make_store_path
from zarr.sync import SyncMixin, sync
from typing import overload

logger = logging.getLogger("zarr.group")

Expand All @@ -41,6 +42,26 @@ def parse_attributes(data: Any) -> dict[str, Any]:
raise TypeError(msg)


@overload
def _parse_async_node(node: AsyncArray) -> Array: ...


@overload
def _parse_async_node(node: AsyncGroup) -> Group: ...


def _parse_async_node(node: AsyncArray | AsyncGroup) -> Array | Group:
"""
Wrap an AsyncArray in an Array, or an AsyncGroup in a Group.
"""
if isinstance(node, AsyncArray):
return Array(node)
elif isinstance(node, Group):
return Group(node)
else:
assert False


@dataclass(frozen=True)
class GroupMetadata(Metadata):
attributes: dict[str, Any] = field(default_factory=dict)
Expand Down Expand Up @@ -509,11 +530,10 @@ def members(self) -> tuple[tuple[str, Array | Group], ...]:
Return the sub-arrays and sub-groups of this group as a tuple of (name, array | group)
pairs
"""
_members: list[AsyncArray | AsyncGroup] = self._sync_iter(self._async_group.members())
return tuple(
(key, Array(value)) if isinstance(value, AsyncArray) else (key, Group(value))
for key, value in _members
)
_members = self._sync_iter(self._async_group.members())

result = tuple(map(lambda kv: (kv[0], _parse_async_node(kv[1])), _members))
return result

def __contains__(self, member) -> bool:
return self._sync(self._async_group.contains(member))
Expand Down

0 comments on commit b996aff

Please sign in to comment.