Refactor/rename v2 metadata fields (#2301)

* remove array v2 metadata data_type and chunk_grid attributes, using dtype and chunks instead * remove inheritance from ArrayMetadata * make asyncarray generic w.r.t metadata * Copy tom augsburger's overload * minimal typedicts for array metadata * clean up imports * make chunk_grid a cached property for v2 metadata * add T suffix to type aliases * make asyncarray.create overloads more specific, and add some comments * use regular names for unions, but use T_ prefix for typevar --------- Co-authored-by: Joe Hamman <[email protected]>
zarr-developers · Oct 10, 2024 · 6b11bb8 · 6b11bb8
1 parent 395604d
commit 6b11bb8
Show file tree

Hide file tree

Showing 8 changed files with 342 additions and 173 deletions.
diff --git a/src/zarr/api/asynchronous.py b/src/zarr/api/asynchronous.py
@@ -12,8 +12,7 @@
 from zarr.core.common import JSON, AccessModeLiteral, ChunkCoords, MemoryOrder, ZarrFormat
 from zarr.core.config import config
 from zarr.core.group import AsyncGroup
-from zarr.core.metadata.v2 import ArrayV2Metadata
-from zarr.core.metadata.v3 import ArrayV3Metadata
+from zarr.core.metadata import ArrayMetadataDict, ArrayV2Metadata, ArrayV3Metadata
 from zarr.errors import NodeTypeValidationError
 from zarr.storage import (
     StoreLike,
@@ -29,7 +28,7 @@
     from zarr.core.chunk_key_encodings import ChunkKeyEncoding
 
     # TODO: this type could use some more thought
-    ArrayLike = AsyncArray | Array | npt.NDArray[Any]
+    ArrayLike = AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | Array | npt.NDArray[Any]
     PathLike = str
 
 __all__ = [
@@ -98,11 +97,11 @@ def _like_args(a: ArrayLike, kwargs: dict[str, Any]) -> dict[str, Any]:
         if isinstance(a.metadata, ArrayV2Metadata):
             new["compressor"] = a.metadata.compressor
             new["filters"] = a.metadata.filters
-
-        if isinstance(a.metadata, ArrayV3Metadata):
-            new["codecs"] = a.metadata.codecs
         else:
-            raise TypeError(f"Unsupported zarr format: {a.metadata.zarr_format}")
+            # TODO: Remove type: ignore statement when type inference improves.
+            # mypy cannot correctly infer the type of a.metadata here for some reason.
+            new["codecs"] = a.metadata.codecs  # type: ignore[unreachable]
+
     else:
         # TODO: set default values compressor/codecs
         # to do this, we may need to evaluate if this is a v2 or v3 array
@@ -199,7 +198,7 @@ async def open(
     path: str | None = None,
     storage_options: dict[str, Any] | None = None,
     **kwargs: Any,  # TODO: type kwargs as valid args to open_array
-) -> AsyncArray | AsyncGroup:
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | AsyncGroup:
     """Convenience function to open a group or array using file-mode-like semantics.
 
     Parameters
@@ -237,11 +236,13 @@ async def open(
     if "shape" not in kwargs and mode in {"a", "w", "w-"}:
         try:
             metadata_dict = await get_array_metadata(store_path, zarr_format=zarr_format)
+            # TODO: remove this cast when we fix typing for array metadata dicts
+            _metadata_dict = cast(ArrayMetadataDict, metadata_dict)
             # for v2, the above would already have raised an exception if not an array
-            zarr_format = metadata_dict["zarr_format"]
-            is_v3_array = zarr_format == 3 and metadata_dict.get("node_type") == "array"
+            zarr_format = _metadata_dict["zarr_format"]
+            is_v3_array = zarr_format == 3 and _metadata_dict.get("node_type") == "array"
             if is_v3_array or zarr_format == 2:
-                return AsyncArray(store_path=store_path, metadata=metadata_dict)
+                return AsyncArray(store_path=store_path, metadata=_metadata_dict)
         except (AssertionError, FileNotFoundError):
             pass
         return await open_group(store=store_path, zarr_format=zarr_format, mode=mode, **kwargs)
@@ -404,7 +405,9 @@ async def tree(*args: Any, **kwargs: Any) -> None:
     raise NotImplementedError
 
 
-async def array(data: npt.ArrayLike, **kwargs: Any) -> AsyncArray:
+async def array(
+    data: npt.ArrayLike, **kwargs: Any
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array filled with `data`.
 
     Parameters
@@ -658,7 +661,7 @@ async def create(
     dimension_names: Iterable[str] | None = None,
     storage_options: dict[str, Any] | None = None,
     **kwargs: Any,
-) -> AsyncArray:
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array.
 
     Parameters
@@ -810,7 +813,9 @@ async def create(
     )
 
 
-async def empty(shape: ChunkCoords, **kwargs: Any) -> AsyncArray:
+async def empty(
+    shape: ChunkCoords, **kwargs: Any
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an empty array.
 
     Parameters
@@ -829,7 +834,9 @@ async def empty(shape: ChunkCoords, **kwargs: Any) -> AsyncArray:
     return await create(shape=shape, fill_value=None, **kwargs)
 
 
-async def empty_like(a: ArrayLike, **kwargs: Any) -> AsyncArray:
+async def empty_like(
+    a: ArrayLike, **kwargs: Any
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an empty array like `a`.
 
     Parameters
@@ -849,7 +856,9 @@ async def empty_like(a: ArrayLike, **kwargs: Any) -> AsyncArray:
 
 
 # TODO: add type annotations for fill_value and kwargs
-async def full(shape: ChunkCoords, fill_value: Any, **kwargs: Any) -> AsyncArray:
+async def full(
+    shape: ChunkCoords, fill_value: Any, **kwargs: Any
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array, with `fill_value` being used as the default value for
     uninitialized portions of the array.
 
@@ -871,7 +880,9 @@ async def full(shape: ChunkCoords, fill_value: Any, **kwargs: Any) -> AsyncArray
 
 
 # TODO: add type annotations for kwargs
-async def full_like(a: ArrayLike, **kwargs: Any) -> AsyncArray:
+async def full_like(
+    a: ArrayLike, **kwargs: Any
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create a filled array like `a`.
 
     Parameters
@@ -892,7 +903,9 @@ async def full_like(a: ArrayLike, **kwargs: Any) -> AsyncArray:
     return await full(**like_kwargs)
 
 
-async def ones(shape: ChunkCoords, **kwargs: Any) -> AsyncArray:
+async def ones(
+    shape: ChunkCoords, **kwargs: Any
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array, with one being used as the default value for
     uninitialized portions of the array.
 
@@ -911,7 +924,9 @@ async def ones(shape: ChunkCoords, **kwargs: Any) -> AsyncArray:
     return await create(shape=shape, fill_value=1, **kwargs)
 
 
-async def ones_like(a: ArrayLike, **kwargs: Any) -> AsyncArray:
+async def ones_like(
+    a: ArrayLike, **kwargs: Any
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array of ones like `a`.
 
     Parameters
@@ -938,7 +953,7 @@ async def open_array(
     path: PathLike | None = None,
     storage_options: dict[str, Any] | None = None,
     **kwargs: Any,  # TODO: type kwargs as valid args to save
-) -> AsyncArray:
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Open an array using file-mode-like semantics.
 
     Parameters
@@ -981,7 +996,9 @@ async def open_array(
         raise
 
 
-async def open_like(a: ArrayLike, path: str, **kwargs: Any) -> AsyncArray:
+async def open_like(
+    a: ArrayLike, path: str, **kwargs: Any
+) -> AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata]:
     """Open a persistent array like `a`.
 
     Parameters
@@ -1004,7 +1021,9 @@ async def open_like(a: ArrayLike, path: str, **kwargs: Any) -> AsyncArray:
     return await open_array(path=path, **like_kwargs)
 
 
-async def zeros(shape: ChunkCoords, **kwargs: Any) -> AsyncArray:
+async def zeros(
+    shape: ChunkCoords, **kwargs: Any
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array, with zero being used as the default value for
     uninitialized portions of the array.
 
@@ -1023,7 +1042,9 @@ async def zeros(shape: ChunkCoords, **kwargs: Any) -> AsyncArray:
     return await create(shape=shape, fill_value=0, **kwargs)
 
 
-async def zeros_like(a: ArrayLike, **kwargs: Any) -> AsyncArray:
+async def zeros_like(
+    a: ArrayLike, **kwargs: Any
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array of zeros like `a`.
 
     Parameters