Skip to content

Commit

Permalink
basic support for v2 and v3 groups
Browse files Browse the repository at this point in the history
  • Loading branch information
jhamman committed Dec 5, 2023
1 parent b4c2a19 commit d52b50c
Show file tree
Hide file tree
Showing 9 changed files with 550 additions and 393 deletions.
10 changes: 5 additions & 5 deletions zarr/v3/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
from zarr.v3.array import Array # noqa: F401
from zarr.v3.array_v2 import ArrayV2 # noqa: F401
from zarr.v3.group import Group # noqa: F401
from zarr.v3.group_v2 import GroupV2 # noqa: F401
from zarr.v3.metadata import RuntimeConfiguration, runtime_configuration # noqa: F401
from zarr.v3.store import ( # noqa: F401
LocalStore,
Expand All @@ -22,18 +21,19 @@
async def open_auto_async(
store: StoreLike,
runtime_configuration_: RuntimeConfiguration = RuntimeConfiguration(),
) -> Union[Array, ArrayV2, Group, GroupV2]:
) -> Union[Array, ArrayV2, Group]:
store_path = make_store_path(store)
try:
return await Group.open_or_array(store_path, runtime_configuration=runtime_configuration_)
return await Array.open(store_path, runtime_configuration=runtime_configuration_)
except KeyError:
return await GroupV2.open_or_array(store_path, runtime_configuration_)
return await Group.open(store_path, runtime_configuration=runtime_configuration_)



def open_auto(
store: StoreLike,
runtime_configuration_: RuntimeConfiguration = RuntimeConfiguration(),
) -> Union[Array, ArrayV2, Group, GroupV2]:
) -> Union[Array, ArrayV2, Group]:
return _sync(
open_auto_async(store, runtime_configuration_),
runtime_configuration_.asyncio_loop,
Expand Down
1 change: 1 addition & 0 deletions zarr/v3/abc/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
from zarr.v3.abc.group import AsyncGroup, SyncGroup
16 changes: 8 additions & 8 deletions zarr/v3/abc/array.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,22 +79,22 @@ def info(self) -> Any:
...


class AsynchronousArray(BaseArray):
class AsyncArray(BaseArray):
"""This class can be implemented as a v2 or v3 array"""

@classmethod
@abstractmethod
async def from_json(cls, zarr_json: Any, store: ReadStore) -> AsynchronousArray:
async def from_json(cls, zarr_json: Any, store: ReadStore) -> AsyncArray:
...

@classmethod
@abstractmethod
async def open(cls, store: ReadStore) -> AsynchronousArray:
async def open(cls, store: ReadStore) -> AsyncArray:
...

@classmethod
@abstractmethod
async def create(cls, store: WriteStore, *, shape, **kwargs) -> AsynchronousArray:
async def create(cls, store: WriteStore, *, shape, **kwargs) -> AsyncArray:
...

@abstractmethod
Expand All @@ -106,24 +106,24 @@ async def setitem(self, selection: Selection, value: np.ndarray) -> None:
...


class SynchronousArray(BaseArray):
class SyncArray(BaseArray):
"""
This class can be implemented as a v2 or v3 array
"""

@classmethod
@abstractmethod
def from_json(cls, zarr_json: Any, store: ReadStore) -> SynchronousArray:
def from_json(cls, zarr_json: Any, store: ReadStore) -> SyncArray:
...

@classmethod
@abstractmethod
def open(cls, store: ReadStore) -> SynchronousArray:
def open(cls, store: ReadStore) -> SyncArray:
...

@classmethod
@abstractmethod
def create(cls, store: WriteStore, *, shape, **kwargs) -> SynchronousArray:
def create(cls, store: WriteStore, *, shape, **kwargs) -> SyncArray:
...

@abstractmethod
Expand Down
250 changes: 185 additions & 65 deletions zarr/v3/abc/group.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
from __future__ import annotations

from abc import abstractproperty, ABC
from abc import abstractproperty, abstractmethod, ABC
from collections.abc import MutableMapping
from typing import Dict, Any
from typing import Dict, Any, AsyncIterator, Union, Iterator

from zarr.v3.abc.array import AsyncArray, SyncArray


class BaseGroup(ABC):
Expand All @@ -17,70 +19,188 @@ def info(self) -> Any: # TODO: type this later
...


class AsynchronousGroup(BaseGroup):
pass
# TODO: (considering the following api)
# store_path (rename to path?)
# nchildren - number of child groups + arrays
# children (async iterator)
# contains - check if child exists
# getitem - get child
# group_keys (async iterator)
# groups (async iterator)
# array_keys (async iterator)
# arrays (async iterator)
# visit
# visitkeys
# visitvalues
# tree
# create_group
# require_group
# create_groups
# require_groups
# create_dataset
# require_dataset
# create
# empty
# zeros
# ones
# full
# array
# empty_like
# zeros_like
# ones_like
# full_like
# move


class SynchronousGroup(BaseGroup, MutableMapping):
# TODO - think about if we want to keep the MutableMapping abstraction or
pass
class AsyncGroup(BaseGroup):
@abstractmethod
async def nchildren(self) -> int:
...

@abstractmethod
async def children(self) -> AsyncIterator:
...

@abstractmethod
async def contains(self, child: str) -> bool:
"""check if child exists"""
...

@abstractmethod
async def getitem(self, child: str) -> Union[AsyncArray, "AsyncGroup"]:
"""get child"""
...

@abstractmethod
async def group_keys(self) -> AsyncIterator[str]:
"""iterate over child group keys"""
...

@abstractmethod
async def groups(self) -> AsyncIterator["AsyncGroup"]:
"""iterate over child groups"""
...

@abstractmethod
async def array_keys(self) -> AsyncIterator[str]:
"""iterate over child array keys"""
...

@abstractmethod
async def arrays(self) -> AsyncIterator[AsyncArray]:
"""iterate over child arrays"""
...

@abstractmethod
async def tree(self, expand=False, level=None) -> Any: # TODO: type this later
...

@abstractmethod
async def create_group(self, name: str, **kwargs) -> "AsyncGroup":
...

@abstractmethod
async def create_array(self, name: str, **kwargs) -> AsyncArray:
...

@abstractmethod
async def empty(self, **kwargs) -> AsyncArray:
...

@abstractmethod
async def zeros(self, **kwargs) -> AsyncArray:
...

@abstractmethod
async def ones(self, **kwargs) -> AsyncArray:
...

@abstractmethod
async def full(self, **kwargs) -> AsyncArray:
...

@abstractmethod
async def empty_like(self, prototype: AsyncArray, **kwargs) -> AsyncArray:
...

@abstractmethod
async def zeros_like(self, prototype: AsyncArray, **kwargs) -> AsyncArray:
...

@abstractmethod
async def ones_like(self, prototype: AsyncArray, **kwargs) -> AsyncArray:
...

@abstractmethod
async def full_like(self, prototype: AsyncArray, **kwargs) -> AsyncArray:
...

@abstractmethod
async def move(self, source: str, dest: str) -> None:
...

# TODO / maybes:
# store_path (rename to path?)
# __enter__
# __exit__
# group_keys
# groups
# array_keys
# arrays
# visit
# visitkeys
# visitvalues
# visititems
# tree
# create_group
# require_group
# create_groups
# require_groups
# create_dataset
# require_dataset
# create
# empty
# zeros
# ones
# full
# array
# empty_like
# zeros_like
# ones_like
# full_like
# move


class SyncGroup(BaseGroup, MutableMapping):
@abstractproperty
def nchildren(self) -> int:
...

@abstractproperty
def children(self) -> Iterator:
...

@abstractmethod
def __contains__(self, child: str) -> bool:
"""check if child exists"""
...

@abstractmethod
def __getitem__(self, child: str) -> Union[SyncArray, "SyncGroup"]:
"""get child"""
...

@abstractmethod
def __setitem__(self, key: str, value: Union[SyncArray, "SyncGroup"]) -> None:
"""get child"""
...

@abstractmethod
def group_keys(self) -> AsyncIterator[str]:
"""iterate over child group keys"""
...

@abstractmethod
def groups(self) -> AsyncIterator["SyncGroup"]:
"""iterate over child groups"""
...

@abstractmethod
def array_keys(self) -> AsyncIterator[str]:
"""iterate over child array keys"""
...

@abstractmethod
def arrays(self) -> AsyncIterator[SyncArray]:
"""iterate over child arrays"""
...

@abstractmethod
def tree(self) -> Any:
...

@abstractmethod
def create_group(self, name: str, **kwargs) -> "SyncGroup":
...

@abstractmethod
def create_array(self, name: str, **kwargs) -> SyncArray:
...

@abstractmethod
def empty(self, **kwargs) -> SyncArray:
...

@abstractmethod
def zeros(self, **kwargs) -> SyncArray:
...

@abstractmethod
def ones(self, **kwargs) -> SyncArray:
...

@abstractmethod
def full(self, **kwargs) -> SyncArray:
...

@abstractmethod
def empty_like(self, prototype: SyncArray, **kwargs) -> SyncArray:
...

@abstractmethod
def zeros_like(self, prototype: SyncArray, **kwargs) -> SyncArray:
...

@abstractmethod
def ones_like(self, prototype: SyncArray, **kwargs) -> SyncArray:
...

@abstractmethod
def full_like(self, prototype: SyncArray, **kwargs) -> SyncArray:
...

@abstractmethod
def move(self, source: str, dest: str) -> None:
...
8 changes: 4 additions & 4 deletions zarr/v3/array.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
# Notes on what I've changed here:
# 1. Split Array into AsyncArray and Array
# 2. Inherit from abc (SynchronousArray, AsynchronousArray)
# 2. Inherit from abc (SyncArray, AsyncArray)
# 3. Added .size and .attrs methods
# 4. Temporarily disabled the creation of ArrayV2
# 5. Added from_json to AsyncArray
Expand All @@ -17,7 +17,7 @@
import numpy as np
from attr import evolve, frozen

from zarr.v3.abc.array import SynchronousArray, AsynchronousArray
from zarr.v3.abc.array import SyncArray, AsyncArray

# from zarr.v3.array_v2 import ArrayV2
from zarr.v3.codecs import CodecMetadata, CodecPipeline, bytes_codec
Expand Down Expand Up @@ -47,7 +47,7 @@


@frozen
class AsyncArray(AsynchronousArray):
class AsyncArray(AsyncArray):
metadata: ArrayMetadata
store_path: StorePath
runtime_configuration: RuntimeConfiguration
Expand Down Expand Up @@ -414,7 +414,7 @@ async def info(self):


@frozen
class Array(SynchronousArray):
class Array(SyncArray):
_async_array: AsyncArray

@classmethod
Expand Down
Loading

0 comments on commit d52b50c

Please sign in to comment.