api: make buffer public api

Author: jhamman

docs/user-guide/config.rst

@@ -70,7 +70,7 @@ This is the current default configuration::
                                         'string': {'name': 'vlen-utf8'}},
               'write_empty_chunks': False},
     'async': {'concurrency': 10, 'timeout': None},
-    'buffer': 'zarr.core.buffer.cpu.Buffer',
+    'buffer': 'zarr.buffer.cpu.Buffer',
     'codec_pipeline': {'batch_size': 1,
                        'path': 'zarr.core.codec_pipeline.BatchedCodecPipeline'},
     'codecs': {'blosc': 'zarr.codecs.blosc.BloscCodec',
@@ -85,5 +85,5 @@ This is the current default configuration::
                'zstd': 'zarr.codecs.zstd.ZstdCodec'},
     'default_zarr_format': 3,
     'json_indent': 2,
-    'ndbuffer': 'zarr.core.buffer.cpu.NDBuffer',
+    'ndbuffer': 'zarr.buffer.cpu.NDBuffer',
     'threading': {'max_workers': None}}

src/zarr/core/buffer/core.py renamed to src/zarr/abc/buffer.py

@@ -23,8 +23,23 @@
     from zarr.codecs.bytes import Endian
     from zarr.core.common import BytesLike, ChunkCoords
 
-# Everything here is imported into ``zarr.core.buffer`` namespace.
-__all__: list[str] = []
+__all__ = [
+    "ArrayLike",
+    "Buffer",
+    "BufferPrototype",
+    "NDArrayLike",
+    "NDBuffer",
+]
+
+
+def _check_item_key_is_1d_contiguous(key: Any) -> None:
+    """Raises error if `key` isn't a 1d contiguous slice"""
+    if not isinstance(key, slice):
+        raise TypeError(
+            f"Item key has incorrect type (expected slice, got {key.__class__.__name__})"
+        )
+    if not (key.step is None or key.step == 1):
+        raise ValueError("slice must be contiguous")
 
 
 @runtime_checkable
@@ -105,16 +120,6 @@ def __eq__(self, other: object) -> Self:  # type: ignore[explicit-override, over
         """
 
 
-def check_item_key_is_1d_contiguous(key: Any) -> None:
-    """Raises error if `key` isn't a 1d contiguous slice"""
-    if not isinstance(key, slice):
-        raise TypeError(
-            f"Item key has incorrect type (expected slice, got {key.__class__.__name__})"
-        )
-    if not (key.step is None or key.step == 1):
-        raise ValueError("slice must be contiguous")
-
-
 class Buffer(ABC):
     """A flat contiguous memory block
 
@@ -266,11 +271,11 @@ def to_bytes(self) -> bytes:
         return bytes(self.as_numpy_array())
 
     def __getitem__(self, key: slice) -> Self:
-        check_item_key_is_1d_contiguous(key)
+        _check_item_key_is_1d_contiguous(key)
         return self.__class__(self._data.__getitem__(key))
 
     def __setitem__(self, key: slice, value: Any) -> None:
-        check_item_key_is_1d_contiguous(key)
+        _check_item_key_is_1d_contiguous(key)
         self._data.__setitem__(key, value)
 
     def __len__(self) -> int:
@@ -498,13 +503,3 @@ class BufferPrototype(NamedTuple):
 
     buffer: type[Buffer]
     nd_buffer: type[NDBuffer]
-
-
-# The default buffer prototype used throughout the Zarr codebase.
-def default_buffer_prototype() -> BufferPrototype:
-    from zarr.registry import (
-        get_buffer_class,
-        get_ndbuffer_class,
-    )
-
-    return BufferPrototype(buffer=get_buffer_class(), nd_buffer=get_ndbuffer_class())

src/zarr/abc/codec.py

@@ -3,8 +3,8 @@
 from abc import abstractmethod
 from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
+from zarr.abc.buffer import Buffer, NDBuffer
 from zarr.abc.metadata import Metadata
-from zarr.core.buffer import Buffer, NDBuffer
 from zarr.core.common import ChunkCoords, concurrent_map
 from zarr.core.config import config
 

src/zarr/abc/store.py

@@ -5,7 +5,7 @@
 from itertools import starmap
 from typing import TYPE_CHECKING, Protocol, runtime_checkable
 
-from zarr.core.buffer.core import default_buffer_prototype
+from zarr.buffer import default_buffer_prototype
 from zarr.core.common import concurrent_map
 from zarr.core.config import config
 
@@ -14,7 +14,7 @@
     from types import TracebackType
     from typing import Any, Self, TypeAlias
 
-    from zarr.core.buffer import Buffer, BufferPrototype
+    from zarr.abc.buffer import Buffer, BufferPrototype
     from zarr.core.common import BytesLike
 
 __all__ = ["ByteGetter", "ByteSetter", "Store", "set_or_delete"]

src/zarr/api/asynchronous.py

@@ -9,9 +9,9 @@
 import numpy.typing as npt
 from typing_extensions import deprecated
 
+from zarr.abc.buffer import NDArrayLike
 from zarr.core.array import Array, AsyncArray, create_array, get_array_metadata
 from zarr.core.array_spec import ArrayConfig, ArrayConfigLike
-from zarr.core.buffer import NDArrayLike
 from zarr.core.common import (
     JSON,
     AccessModeLiteral,

src/zarr/api/synchronous.py

@@ -16,6 +16,7 @@
 
     import numpy.typing as npt
 
+    from zarr.abc.buffer import NDArrayLike
     from zarr.abc.codec import Codec
     from zarr.api.asynchronous import ArrayLike, PathLike
     from zarr.core.array import (
@@ -25,7 +26,6 @@
         ShardsLike,
     )
     from zarr.core.array_spec import ArrayConfig, ArrayConfigLike
-    from zarr.core.buffer import NDArrayLike
     from zarr.core.chunk_key_encodings import ChunkKeyEncoding, ChunkKeyEncodingLike
     from zarr.core.common import (
         JSON,

src/zarr/buffer/__init__.py

@@ -0,0 +1,17 @@
+from zarr.abc.buffer import BufferPrototype
+from zarr.buffer.cpu import numpy_buffer_prototype
+
+__all__ = [
+    "default_buffer_prototype",
+    "numpy_buffer_prototype",
+]
+
+
+# The default buffer prototype used throughout the Zarr codebase.
+def default_buffer_prototype() -> BufferPrototype:
+    from zarr.registry import (
+        get_buffer_class,
+        get_ndbuffer_class,
+    )
+
+    return BufferPrototype(buffer=get_buffer_class(), nd_buffer=get_ndbuffer_class())

src/zarr/core/buffer/cpu.py renamed to src/zarr/buffer/cpu.py

@@ -9,7 +9,7 @@
 import numpy as np
 import numpy.typing as npt
 
-from zarr.core.buffer import core
+import zarr.abc.buffer
 from zarr.registry import (
     register_buffer,
     register_ndbuffer,
@@ -19,11 +19,11 @@
     from collections.abc import Callable, Iterable
     from typing import Self
 
-    from zarr.core.buffer.core import ArrayLike, NDArrayLike
+    from zarr.abc.buffer import ArrayLike, NDArrayLike
     from zarr.core.common import BytesLike
 
 
-class Buffer(core.Buffer):
+class Buffer(zarr.abc.buffer.Buffer):
     """A flat contiguous memory block
 
     We use Buffer throughout Zarr to represent a contiguous block of memory.
@@ -52,7 +52,7 @@ def create_zero_length(cls) -> Self:
         return cls(np.array([], dtype="b"))
 
     @classmethod
-    def from_buffer(cls, buffer: core.Buffer) -> Self:
+    def from_buffer(cls, buffer: zarr.abc.buffer.Buffer) -> Self:
         """Create a new buffer of an existing Buffer
 
         This is useful if you want to ensure that an existing buffer is
@@ -107,7 +107,7 @@ def as_numpy_array(self) -> npt.NDArray[Any]:
         """
         return np.asanyarray(self._data)
 
-    def __add__(self, other: core.Buffer) -> Self:
+    def __add__(self, other: zarr.abc.buffer.Buffer) -> Self:
         """Concatenate two buffers"""
 
         other_array = other.as_array_like()
@@ -117,7 +117,7 @@ def __add__(self, other: core.Buffer) -> Self:
         )
 
 
-class NDBuffer(core.NDBuffer):
+class NDBuffer(zarr.abc.buffer.NDBuffer):
     """An n-dimensional memory block
 
     We use NDBuffer throughout Zarr to represent a n-dimensional memory block.
@@ -186,8 +186,10 @@ def __setitem__(self, key: Any, value: Any) -> None:
 
 
 def as_numpy_array_wrapper(
-    func: Callable[[npt.NDArray[Any]], bytes], buf: core.Buffer, prototype: core.BufferPrototype
-) -> core.Buffer:
+    func: Callable[[npt.NDArray[Any]], bytes],
+    buf: zarr.abc.buffer.Buffer,
+    prototype: zarr.abc.buffer.BufferPrototype,
+) -> zarr.abc.buffer.Buffer:
     """Converts the input of `func` to a numpy array and the output back to `Buffer`.
 
     This function is useful when calling a `func` that only support host memory such
@@ -214,13 +216,13 @@ def as_numpy_array_wrapper(
 
 
 # CPU buffer prototype using numpy arrays
-buffer_prototype = core.BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
+buffer_prototype = zarr.abc.buffer.BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
 # default_buffer_prototype = buffer_prototype
 
 
 # The numpy prototype used for E.g. when reading the shard index
-def numpy_buffer_prototype() -> core.BufferPrototype:
-    return core.BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
+def numpy_buffer_prototype() -> zarr.abc.buffer.BufferPrototype:
+    return zarr.abc.buffer.BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
 
 
 register_buffer(Buffer)

src/zarr/core/buffer/gpu.py renamed to src/zarr/buffer/gpu.py

@@ -11,8 +11,7 @@
 import numpy as np
 import numpy.typing as npt
 
-from zarr.core.buffer import core
-from zarr.core.buffer.core import ArrayLike, BufferPrototype, NDArrayLike
+import zarr.abc.buffer
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -26,7 +25,7 @@
     cp = None
 
 
-class Buffer(core.Buffer):
+class Buffer(zarr.abc.buffer.Buffer):
     """A flat contiguous memory block on the GPU
 
     We use Buffer throughout Zarr to represent a contiguous block of memory.
@@ -47,7 +46,7 @@ class Buffer(core.Buffer):
         array-like object that must be 1-dim, contiguous, and byte dtype.
     """
 
-    def __init__(self, array_like: ArrayLike) -> None:
+    def __init__(self, array_like: zarr.abc.buffer.ArrayLike) -> None:
         if cp is None:
             raise ImportError(
                 "Cannot use zarr.buffer.gpu.Buffer without cupy. Please install cupy."
@@ -83,7 +82,7 @@ def create_zero_length(cls) -> Self:
         return cls(cp.array([], dtype="b"))
 
     @classmethod
-    def from_buffer(cls, buffer: core.Buffer) -> Self:
+    def from_buffer(cls, buffer: zarr.abc.buffer.Buffer) -> Self:
         """Create an GPU Buffer given an arbitrary Buffer
         This will try to be zero-copy if `buffer` is already on the
         GPU and will trigger a copy if not.
@@ -101,7 +100,7 @@ def from_bytes(cls, bytes_like: BytesLike) -> Self:
     def as_numpy_array(self) -> npt.NDArray[Any]:
         return cast(npt.NDArray[Any], cp.asnumpy(self._data))
 
-    def __add__(self, other: core.Buffer) -> Self:
+    def __add__(self, other: zarr.abc.buffer.Buffer) -> Self:
         other_array = other.as_array_like()
         assert other_array.dtype == np.dtype("b")
         gpu_other = Buffer(other_array)
@@ -111,7 +110,7 @@ def __add__(self, other: core.Buffer) -> Self:
         )
 
 
-class NDBuffer(core.NDBuffer):
+class NDBuffer(zarr.abc.buffer.NDBuffer):
     """A n-dimensional memory block on the GPU
 
     We use NDBuffer throughout Zarr to represent a n-dimensional memory block.
@@ -136,7 +135,7 @@ class NDBuffer(core.NDBuffer):
         ndarray-like object that is convertible to a regular Numpy array.
     """
 
-    def __init__(self, array: NDArrayLike) -> None:
+    def __init__(self, array: zarr.abc.buffer.NDArrayLike) -> None:
         if cp is None:
             raise ImportError(
                 "Cannot use zarr.buffer.gpu.NDBuffer without cupy. Please install cupy."
@@ -208,10 +207,10 @@ def __getitem__(self, key: Any) -> Self:
     def __setitem__(self, key: Any, value: Any) -> None:
         if isinstance(value, NDBuffer):
             value = value._data
-        elif isinstance(value, core.NDBuffer):
+        elif isinstance(value, zarr.abc.buffer.NDBuffer):
             gpu_value = NDBuffer(value.as_ndarray_like())
             value = gpu_value._data
         self._data.__setitem__(key, value)
 
 
-buffer_prototype = BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
+buffer_prototype = zarr.abc.buffer.BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)

src/zarr/codecs/blosc.py

@@ -10,15 +10,15 @@
 from numcodecs.blosc import Blosc
 
 from zarr.abc.codec import BytesBytesCodec
-from zarr.core.buffer.cpu import as_numpy_array_wrapper
+from zarr.buffer.cpu import as_numpy_array_wrapper
 from zarr.core.common import JSON, parse_enum, parse_named_configuration
 from zarr.registry import register_codec
 
 if TYPE_CHECKING:
     from typing import Self
 
+    from zarr.abc.buffer import Buffer
     from zarr.core.array_spec import ArraySpec
-    from zarr.core.buffer import Buffer
 
 
 class BloscShuffle(Enum):