add ArrayModel.from_array, await unawaited awaitables, add flattening tests

Author: d-v-b

src/zarr/group.py

@@ -183,10 +183,10 @@ async def open(
             assert zarr_json_bytes is not None
             group_metadata = json.loads(zarr_json_bytes.to_bytes())
 
-        return cls.from_dict(store_path, group_metadata)
+        return await cls.from_dict(store_path, group_metadata)
 
     @classmethod
-    def from_dict(
+    async def from_dict(
         cls,
         store_path: StorePath,
         data: dict[str, Any],
@@ -217,9 +217,9 @@ async def getitem(
             else:
                 zarr_json = json.loads(zarr_json_bytes.to_bytes())
             if zarr_json["node_type"] == "group":
-                return type(self).from_dict(store_path, zarr_json)
+                return await type(self).from_dict(store_path, zarr_json)
             elif zarr_json["node_type"] == "array":
-                return sync(AsyncArray.from_dict(store_path, zarr_json))
+                return await AsyncArray.from_dict(store_path, zarr_json)
             else:
                 raise ValueError(f"unexpected node_type: {zarr_json['node_type']}")
         elif self.metadata.zarr_format == 2:
@@ -250,7 +250,7 @@ async def getitem(
                     else {"zarr_format": self.metadata.zarr_format}
                 )
                 zarr_json = {**zgroup, "attributes": zattrs}
-                return type(self).from_dict(store_path, zarr_json)
+                return await type(self).from_dict(store_path, zarr_json)
         else:
             raise ValueError(f"unexpected zarr_format: {self.metadata.zarr_format}")
 

src/zarr/hierarchy.py

@@ -13,13 +13,111 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+from typing import Any, Literal
 
+import numpy as np
 from typing_extensions import Self
 
+from zarr.abc.codec import CodecPipeline
 from zarr.array import Array
+from zarr.buffer import NDBuffer
+from zarr.chunk_grids import ChunkGrid, RegularChunkGrid
+from zarr.chunk_key_encodings import ChunkKeyEncoding, DefaultChunkKeyEncoding
+from zarr.codecs.bytes import BytesCodec
 from zarr.group import Group, GroupMetadata
 from zarr.metadata import ArrayV3Metadata
 from zarr.store.core import StorePath
+from zarr.v2.util import guess_chunks
+
+
+def auto_data_type(data: Any) -> Any:
+    if hasattr(data, "dtype"):
+        if hasattr(data, "data_type"):
+            msg = (
+                f"Could not infer the data_type attribute from {data}, because "
+                "it has both `dtype` and `data_type` attributes. "
+                "This method requires input with one, or the other, of these attributes."
+            )
+            raise ValueError(msg)
+        return data.dtype
+    elif hasattr(data, "data_type") and not hasattr(data, "dtype"):
+        return data.data_type
+    else:
+        msg = (
+            f"Could not infer the data_type attribute from {data}. "
+            "Expected either an object with a `dtype` attribute, "
+            "or an object with a `data_type` attribute."
+        )
+        raise ValueError(msg)
+
+
+def auto_attributes(data: Any) -> Any:
+    """
+    Guess attributes from:
+        input with an `attrs` attribute, or
+        input with an `attributes` attribute,
+        or anything (returning {})
+    """
+    if hasattr(data, "attrs"):
+        return data.attrs
+    if hasattr(data, "attributes"):
+        return data.attributes
+    return {}
+
+
+def auto_chunk_key_encoding(data: Any) -> Any:
+    if hasattr(data, "chunk_key_encoding"):
+        return data.chunk_key_encoding
+    return DefaultChunkKeyEncoding()
+
+
+def auto_fill_value(data: Any) -> Any:
+    """
+    Guess fill value from an input with a `fill_value` attribute, returning 0 otherwise.
+    """
+    if hasattr(data, "fill_value"):
+        return data.fill_value
+    return 0
+
+
+def auto_codecs(data: Any) -> Any:
+    """
+    Guess compressor from an input with a `compressor` attribute, returning `None` otherwise.
+    """
+    if hasattr(data, "codecs"):
+        return data.codecs
+    return (BytesCodec(),)
+
+
+def auto_dimension_names(data: Any) -> Any:
+    """
+    If the input has a `dimension_names` attribute, return it, otherwise
+    return None.
+    """
+
+    if hasattr(data, "dimension_names"):
+        return data.dimension_names
+    return None
+
+
+def auto_chunk_grid(data: Any) -> Any:
+    """
+    Guess a chunk grid from:
+      input with a `chunk_grid` attribute,
+      input with a `chunksize` attribute, or
+      input with a `chunks` attribute, or,
+      input with `shape` and `dtype` attributes
+    """
+    if hasattr(data, "chunk_grid"):
+        # more a statement of intent than anything else
+        return data.chunk_grid
+    if hasattr(data, "chunksize"):
+        chunks = data.chunksize
+    elif hasattr(data, "chunks"):
+        chunks = data.chunks
+    else:
+        chunks = guess_chunks(data.shape, np.dtype(data.dtype).itemsize)
+    return RegularChunkGrid(chunk_shape=chunks)
 
 
 class ArrayModel(ArrayV3Metadata):
@@ -43,6 +141,70 @@ def to_stored(self, store_path: StorePath, exists_ok: bool = False) -> Array:
 
         return Array.from_dict(store_path=store_path, data=self.to_dict())
 
+    @classmethod
+    def from_array(
+        cls: type[Self],
+        data: NDBuffer,
+        *,
+        chunk_grid: ChunkGrid | Literal["auto"] = "auto",
+        chunk_key_encoding: ChunkKeyEncoding | Literal["auto"] = "auto",
+        fill_value: Any | Literal["auto"] = "auto",
+        codecs: CodecPipeline | Literal["auto"] = "auto",
+        attributes: dict[str, Any] | Literal["auto"] = "auto",
+        dimension_names: tuple[str, ...] | Literal["auto"] = "auto",
+    ) -> Self:
+        """
+        Create an ArrayModel from an array-like object, e.g. a numpy array.
+
+        The returned ArrayModel will use the shape and dtype attributes of the input.
+        The remaining ArrayModel attributes are exposed by this method as keyword arguments,
+        which can either be the string "auto", which instructs this method to infer or guess
+        a value, or a concrete value to use.
+        """
+        shape_out = data.shape
+        data_type_out = auto_data_type(data)
+
+        if chunk_grid == "auto":
+            chunk_grid_out = auto_chunk_grid(data)
+        else:
+            chunk_grid_out = chunk_grid
+
+        if chunk_key_encoding == "auto":
+            chunk_key_encoding_out = auto_chunk_key_encoding(data)
+        else:
+            chunk_key_encoding_out = chunk_key_encoding
+
+        if fill_value == "auto":
+            fill_value_out = auto_fill_value(data)
+        else:
+            fill_value_out = fill_value
+
+        if codecs == "auto":
+            codecs_out = auto_codecs(data)
+        else:
+            codecs_out = codecs
+
+        if attributes == "auto":
+            attributes_out = auto_attributes(data)
+        else:
+            attributes_out = attributes
+
+        if dimension_names == "auto":
+            dimension_names_out = auto_dimension_names(data)
+        else:
+            dimension_names_out = dimension_names
+
+        return cls(
+            shape=shape_out,
+            data_type=data_type_out,
+            chunk_grid=chunk_grid_out,
+            chunk_key_encoding=chunk_key_encoding_out,
+            fill_value=fill_value_out,
+            codecs=codecs_out,
+            attributes=attributes_out,
+            dimension_names=dimension_names_out,
+        )
+
 
 @dataclass(frozen=True)
 class GroupModel(GroupMetadata):
@@ -104,11 +266,20 @@ def to_flat(
     """
     result = {}
     model_copy: ArrayModel | GroupModel
-    node_dict = node.to_dict()
     if isinstance(node, ArrayModel):
-        model_copy = ArrayModel(**node_dict)
+        # we can remove this if we add a model_copy method
+        model_copy = ArrayModel(
+            shape=node.shape,
+            data_type=node.data_type,
+            chunk_grid=node.chunk_grid,
+            chunk_key_encoding=node.chunk_key_encoding,
+            fill_value=node.fill_value,
+            codecs=node.codecs,
+            attributes=node.attributes,
+            dimension_names=node.dimension_names,
+        )
     else:
-        model_copy = GroupModel(**node_dict)
+        model_copy = GroupModel(attributes=node.attributes, members=None)
         if node.members is not None:
             for name, value in node.members.items():
                 result.update(to_flat(value, "/".join([root_path, name])))

src/zarr/metadata.py

@@ -162,7 +162,7 @@ class ArrayV3Metadata(ArrayMetadata):
     chunk_key_encoding: ChunkKeyEncoding
     fill_value: Any
     codecs: CodecPipeline
-    attributes: dict[str, Any] = field(default_factory=dict)
+    attributes: dict[str, JSON] = field(default_factory=dict)
     dimension_names: tuple[str, ...] | None = None
     zarr_format: Literal[3] = field(default=3, init=False)
     node_type: Literal["array"] = field(default="array", init=False)

tests/v3/test_hierarchy.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
+import numpy as np
 import pytest
 
 from zarr.array import Array
 from zarr.chunk_grids import RegularChunkGrid
 from zarr.chunk_key_encodings import DefaultChunkKeyEncoding
 from zarr.group import GroupMetadata
-from zarr.hierarchy import ArrayModel, GroupModel
+from zarr.hierarchy import ArrayModel, GroupModel, from_flat, to_flat
 from zarr.metadata import ArrayV3Metadata
 from zarr.store.core import StorePath
 from zarr.store.memory import MemoryStore
@@ -36,7 +37,7 @@ def test_array_model_to_stored(memory_store: MemoryStore) -> None:
         attributes={"foo": 10},
     )
 
-    array = model.to_stored(memory_store)
+    array = model.to_stored(store_path=StorePath(store=memory_store))
     assert array.metadata.to_dict() == model.to_dict()
 
 
@@ -50,7 +51,7 @@ def test_array_model_from_stored(memory_store: MemoryStore) -> None:
         attributes={"foo": 10},
     )
 
-    array = Array.from_dict(memory_store, array_meta.to_dict())
+    array = Array.from_dict(StorePath(memory_store), array_meta.to_dict())
     array_model = ArrayModel.from_stored(array)
     assert array_model.to_dict() == array_meta.to_dict()
 
@@ -105,3 +106,46 @@ def test_groupmodel_to_stored(
         assert model_rt.members == model.members
     else:
         assert model_rt.members == {}
+
+
+@pytest.mark.parametrize(
+    ("data, expected"),
+    [
+        (
+            ArrayModel.from_array(np.arange(10)),
+            {"": ArrayModel.from_array(np.arange(10))},
+        ),
+        (
+            GroupModel(
+                attributes={"foo": 10},
+                members={"a": ArrayModel.from_array(np.arange(5), attributes={"foo": 100})},
+            ),
+            {
+                "": GroupModel(attributes={"foo": 10}, members=None),
+                "/a": ArrayModel.from_array(np.arange(5), attributes={"foo": 100}),
+            },
+        ),
+        (
+            GroupModel(
+                attributes={},
+                members={
+                    "a": GroupModel(
+                        attributes={"foo": 10},
+                        members={"a": ArrayModel.from_array(np.arange(5), attributes={"foo": 100})},
+                    ),
+                    "b": ArrayModel.from_array(np.arange(2), attributes={"foo": 3}),
+                },
+            ),
+            {
+                "": GroupModel(attributes={}, members=None),
+                "/a": GroupModel(attributes={"foo": 10}, members=None),
+                "/a/a": ArrayModel.from_array(np.arange(5), attributes={"foo": 100}),
+                "/b": ArrayModel.from_array(np.arange(2), attributes={"foo": 3}),
+            },
+        ),
+    ],
+)
+def test_flatten_unflatten(data, expected) -> None:
+    flattened = to_flat(data)
+    assert flattened == expected
+    assert from_flat(flattened) == data