Merge branch 'main' into rmg/ArrayNotFoundError

Author: d-v-b

changes/3318.misc.rst

@@ -0,0 +1,2 @@
+Define a ``Protocol`` to model the ``numcodecs.abc.Codec`` interface. This is groundwork toward
+making ``numcodecs`` an optional dependency for ``zarr-python``.

changes/3371.misc.rst

@@ -0,0 +1 @@
+Ensure that tests for executable examples are run in a fresh python environment.

changes/3372.misc.rst

@@ -0,0 +1,2 @@
+Make certain imports in ``zarr.abc.store`` local to method definitions. This minimizes the risk of
+circular imports when adding new classes to ``zarr.abc.store``.

changes/3374.misc.rst

@@ -0,0 +1 @@
+Replaces usage of the ``zarr.core.common.ChunkCoords`` typealias with ``tuple[int, ...]``.

changes/3378.bugfix.rst

@@ -0,0 +1,3 @@
+Ensure passing `config` is handled properly when `open`ing an existing
+array.
+

docs/developers/roadmap.rst

@@ -56,7 +56,7 @@ Goals
 -  Provide a complete implementation of Zarr V3 through the Zarr-Python
    API
 -  Clear the way for exciting extensions / ZEPs
-   (i.e. `sharding <https://zarr-specs.readthedocs.io/en/latest/v3/codecs/sharding-indexed/v1.0.html>`__,
+   (i.e. `sharding <https://zarr-specs.readthedocs.io/en/latest/v3/codecs/sharding-indexed/>`__,
    `variable chunking <https://zarr.dev/zeps/draft/ZEP0003.html>`__,
    etc.)
 -  Provide a developer API that can be used to implement and register V3

src/zarr/abc/codec.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Generic, TypeVar
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Generic, TypeGuard, TypeVar
+
+from typing_extensions import ReadOnly, TypedDict
 
 from zarr.abc.metadata import Metadata
 from zarr.core.buffer import Buffer, NDBuffer
-from zarr.core.common import ChunkCoords, concurrent_map
+from zarr.core.common import NamedConfig, concurrent_map
 from zarr.core.config import config
 
 if TYPE_CHECKING:
@@ -34,6 +37,27 @@
 CodecInput = TypeVar("CodecInput", bound=NDBuffer | Buffer)
 CodecOutput = TypeVar("CodecOutput", bound=NDBuffer | Buffer)
 
+TName = TypeVar("TName", bound=str, covariant=True)
+
+
+class CodecJSON_V2(TypedDict, Generic[TName]):
+    """The JSON representation of a codec for Zarr V2"""
+
+    id: ReadOnly[TName]
+
+
+def _check_codecjson_v2(data: object) -> TypeGuard[CodecJSON_V2[str]]:
+    return isinstance(data, Mapping) and "id" in data and isinstance(data["id"], str)
+
+
+CodecJSON_V3 = str | NamedConfig[str, Mapping[str, object]]
+"""The JSON representation of a codec for Zarr V3."""
+
+# The widest type we will *accept* for a codec JSON
+# This covers v2 and v3
+CodecJSON = str | Mapping[str, object]
+"""The widest type of JSON-like input that could specify a codec."""
+
 
 class BaseCodec(Metadata, Generic[CodecInput, CodecOutput]):
     """Generic base class for codecs.
@@ -96,7 +120,7 @@ def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
     def validate(
         self,
         *,
-        shape: ChunkCoords,
+        shape: tuple[int, ...],
         dtype: ZDType[TBaseDType, TBaseScalar],
         chunk_grid: ChunkGrid,
     ) -> None:
@@ -105,7 +129,7 @@ def validate(
 
         Parameters
         ----------
-        shape : ChunkCoords
+        shape : tuple[int, ...]
             The array shape
         dtype : np.dtype[Any]
             The array data type
@@ -311,14 +335,18 @@ def supports_partial_encode(self) -> bool: ...
 
     @abstractmethod
     def validate(
-        self, *, shape: ChunkCoords, dtype: ZDType[TBaseDType, TBaseScalar], chunk_grid: ChunkGrid
+        self,
+        *,
+        shape: tuple[int, ...],
+        dtype: ZDType[TBaseDType, TBaseScalar],
+        chunk_grid: ChunkGrid,
     ) -> None:
         """Validates that all codec configurations are compatible with the array metadata.
         Raises errors when a codec configuration is not compatible.
 
         Parameters
         ----------
-        shape : ChunkCoords
+        shape : tuple[int, ...]
             The array shape
         dtype : np.dtype[Any]
             The array data type

src/zarr/abc/numcodec.py

@@ -0,0 +1,101 @@
+from typing import Any, Self, TypeGuard
+
+from typing_extensions import Protocol
+
+
+class Numcodec(Protocol):
+    """
+    A protocol that models the ``numcodecs.abc.Codec`` interface.
+
+    This protocol should be considered experimental. Expect the type annotations for ``buf`` and
+    ``out`` to narrow in the future.
+    """
+
+    codec_id: str
+
+    def encode(self, buf: Any) -> Any:
+        """Encode data from ``buf``.
+
+        Parameters
+        ----------
+        buf : Any
+            Data to be encoded.
+
+        Returns
+        -------
+        enc: Any
+            Encoded data.
+        """
+        ...
+
+    def decode(self, buf: Any, out: Any | None = None) -> Any:
+        """
+        Decode data in ``buf``.
+
+        Parameters
+        ----------
+        buf : Any
+            Encoded data.
+        out : Any
+            Writeable buffer to store decoded data. If provided, this buffer must
+            be exactly the right size to store the decoded data.
+
+        Returns
+        -------
+        dec : Any
+            Decoded data.
+        """
+        ...
+
+    def get_config(self) -> Any:
+        """
+        Return a JSON-serializable configuration dictionary for this
+        codec. Must include an ``'id'`` field with the codec identifier.
+        """
+        ...
+
+    @classmethod
+    def from_config(cls, config: Any) -> Self:
+        """
+        Instantiate a codec from a configuration dictionary.
+
+        Parameters
+        ----------
+        config : Any
+            A configuration dictionary for this codec.
+        """
+        ...
+
+
+def _is_numcodec_cls(obj: object) -> TypeGuard[type[Numcodec]]:
+    """
+    Check if the given object is a class implements the Numcodec protocol.
+
+    The @runtime_checkable decorator does not allow issubclass checks for protocols with non-method
+    members (i.e., attributes), so we use this function to manually check for the presence of the
+    required attributes and methods on a given object.
+    """
+    return (
+        isinstance(obj, type)
+        and hasattr(obj, "codec_id")
+        and isinstance(obj.codec_id, str)
+        and hasattr(obj, "encode")
+        and callable(obj.encode)
+        and hasattr(obj, "decode")
+        and callable(obj.decode)
+        and hasattr(obj, "get_config")
+        and callable(obj.get_config)
+        and hasattr(obj, "from_config")
+        and callable(obj.from_config)
+    )
+
+
+def _is_numcodec(obj: object) -> TypeGuard[Numcodec]:
+    """
+    Check if the given object implements the Numcodec protocol.
+
+    The @runtime_checkable decorator does not allow issubclass checks for protocols with non-method
+    members (i.e., attributes), so we use this function to manually check for the presence of the
+    required attributes and methods on a given object.
+    """
+    return _is_numcodec_cls(type(obj))

src/zarr/abc/store.py

@@ -6,10 +6,6 @@
 from itertools import starmap
 from typing import TYPE_CHECKING, Protocol, runtime_checkable
 
-from zarr.core.buffer.core import default_buffer_prototype
-from zarr.core.common import concurrent_map
-from zarr.core.config import config
-
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, AsyncIterator, Iterable
     from types import TracebackType
@@ -438,6 +434,9 @@ async def getsize(self, key: str) -> int:
         # Note to implementers: this default implementation is very inefficient since
         # it requires reading the entire object. Many systems will have ways to get the
         # size of an object without reading it.
+        # avoid circular import
+        from zarr.core.buffer.core import default_buffer_prototype
+
         value = await self.get(key, prototype=default_buffer_prototype())
         if value is None:
             raise FileNotFoundError(key)
@@ -476,6 +475,11 @@ async def getsize_prefix(self, prefix: str) -> int:
         # on to getting sizes. Ideally we would overlap those two, which should
         # improve tail latency and might reduce memory pressure (since not all keys
         # would be in memory at once).
+
+        # avoid circular import
+        from zarr.core.common import concurrent_map
+        from zarr.core.config import config
+
         keys = [(x,) async for x in self.list_prefix(prefix)]
         limit = config.get("async.concurrency")
         sizes = await concurrent_map(keys, self.getsize, limit=limit)

src/zarr/api/asynchronous.py

@@ -24,7 +24,6 @@
 from zarr.core.common import (
     JSON,
     AccessModeLiteral,
-    ChunkCoords,
     DimensionNames,
     MemoryOrder,
     ZarrFormat,
@@ -53,9 +52,8 @@
 if TYPE_CHECKING:
     from collections.abc import Iterable
 
-    import numcodecs.abc
-
     from zarr.abc.codec import Codec
+    from zarr.abc.numcodec import Numcodec
     from zarr.core.buffer import NDArrayLikeOrScalar
     from zarr.core.chunk_key_encodings import ChunkKeyEncoding
     from zarr.storage import StoreLike
@@ -108,7 +106,7 @@ def _infer_overwrite(mode: AccessModeLiteral) -> bool:
     return mode in _OVERWRITE_MODES
 
 
-def _get_shape_chunks(a: ArrayLike | Any) -> tuple[ChunkCoords | None, ChunkCoords | None]:
+def _get_shape_chunks(a: ArrayLike | Any) -> tuple[tuple[int, ...] | None, tuple[int, ...] | None]:
     """Helper function to get the shape and chunks from an array-like object"""
     shape = None
     chunks = None
@@ -360,7 +358,9 @@ async def open(
             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, config=kwargs.get("config")
+                )
         except (AssertionError, FileNotFoundError, NodeTypeValidationError):
             pass
         return await open_group(store=store_path, zarr_format=zarr_format, mode=mode, **kwargs)
@@ -866,9 +866,9 @@ async def open_group(
 
 
 async def create(
-    shape: ChunkCoords | int,
+    shape: tuple[int, ...] | int,
     *,  # Note: this is a change from v2
-    chunks: ChunkCoords | int | bool | None = None,
+    chunks: tuple[int, ...] | int | bool | None = None,
     dtype: ZDTypeLike | None = None,
     compressor: CompressorLike = "auto",
     fill_value: Any | None = DEFAULT_FILL_VALUE,
@@ -878,7 +878,7 @@ async def create(
     overwrite: bool = False,
     path: PathLike | None = None,
     chunk_store: StoreLike | None = None,
-    filters: Iterable[dict[str, JSON] | numcodecs.abc.Codec] | None = None,
+    filters: Iterable[dict[str, JSON] | Numcodec] | None = None,
     cache_metadata: bool | None = None,
     cache_attrs: bool | None = None,
     read_only: bool | None = None,
@@ -890,7 +890,7 @@ async def create(
     meta_array: Any | None = None,  # TODO: need type
     attributes: dict[str, JSON] | None = None,
     # v3 only
-    chunk_shape: ChunkCoords | int | None = None,
+    chunk_shape: tuple[int, ...] | int | None = None,
     chunk_key_encoding: (
         ChunkKeyEncoding
         | tuple[Literal["default"], Literal[".", "/"]]
@@ -1075,7 +1075,7 @@ async def create(
 
 
 async def empty(
-    shape: ChunkCoords, **kwargs: Any
+    shape: tuple[int, ...], **kwargs: Any
 ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an empty array with the specified shape. The contents will be filled with the
     array's fill value or zeros if no fill value is provided.
@@ -1127,7 +1127,7 @@ async def empty_like(
 
 # TODO: add type annotations for fill_value and kwargs
 async def full(
-    shape: ChunkCoords, fill_value: Any, **kwargs: Any
+    shape: tuple[int, ...], 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.
@@ -1174,7 +1174,7 @@ async def full_like(
 
 
 async def ones(
-    shape: ChunkCoords, **kwargs: Any
+    shape: tuple[int, ...], **kwargs: Any
 ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array, with one being used as the default value for
     uninitialized portions of the array.
@@ -1297,7 +1297,7 @@ async def open_like(
 
 
 async def zeros(
-    shape: ChunkCoords, **kwargs: Any
+    shape: tuple[int, ...], **kwargs: Any
 ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array, with zero being used as the default value for
     uninitialized portions of the array.