Clean up zarr.core.buffer API

Author: dstansby

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.core.buffer._core import default_buffer_prototype
 from zarr.core.common import concurrent_map
 from zarr.core.config import config
 

src/zarr/core/buffer/__init__.py

@@ -1,4 +1,4 @@
-from zarr.core.buffer.core import (
+from zarr.core.buffer._core import (
     ArrayLike,
     Buffer,
     BufferPrototype,

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

@@ -23,9 +23,6 @@
     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] = []
-
 
 @runtime_checkable
 class ArrayLike(Protocol):

src/zarr/core/buffer/cpu.py

@@ -9,7 +9,7 @@
 import numpy as np
 import numpy.typing as npt
 
-from zarr.core.buffer import core
+import zarr.core.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.core.buffer import ArrayLike, NDArrayLike
     from zarr.core.common import BytesLike
 
 
-class Buffer(core.Buffer):
+class Buffer(zarr.core.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.core.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.core.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.core.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.core.buffer.Buffer,
+    prototype: zarr.core.buffer.BufferPrototype,
+) -> zarr.core.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.core.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.core.buffer.BufferPrototype:
+    return zarr.core.buffer.BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
 
 
 register_buffer(Buffer)

src/zarr/core/buffer/gpu.py

@@ -11,8 +11,8 @@
 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.core.buffer
+from zarr.core.buffer import ArrayLike, BufferPrototype, NDArrayLike
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -26,7 +26,7 @@
     cp = None
 
 
-class Buffer(core.Buffer):
+class Buffer(zarr.core.buffer.Buffer):
     """A flat contiguous memory block on the GPU
 
     We use Buffer throughout Zarr to represent a contiguous block of memory.
@@ -83,7 +83,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.core.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 +101,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.core.buffer.Buffer) -> Self:
         other_array = other.as_array_like()
         assert other_array.dtype == np.dtype("b")
         gpu_other = Buffer(other_array)
@@ -111,7 +111,7 @@ def __add__(self, other: core.Buffer) -> Self:
         )
 
 
-class NDBuffer(core.NDBuffer):
+class NDBuffer(zarr.core.buffer.NDBuffer):
     """A n-dimensional memory block on the GPU
 
     We use NDBuffer throughout Zarr to represent a n-dimensional memory block.
@@ -208,7 +208,7 @@ 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.core.buffer.NDBuffer):
             gpu_value = NDBuffer(value.as_ndarray_like())
             value = gpu_value._data
         self._data.__setitem__(key, value)

src/zarr/core/metadata/v3.py

@@ -4,7 +4,7 @@
 from typing import TYPE_CHECKING, TypedDict, overload
 
 from zarr.abc.metadata import Metadata
-from zarr.core.buffer.core import default_buffer_prototype
+from zarr.core.buffer._core import default_buffer_prototype
 
 if TYPE_CHECKING:
     from typing import Self

src/zarr/storage/_local.py

@@ -9,7 +9,7 @@
 
 from zarr.abc.store import ByteRangeRequest, Store
 from zarr.core.buffer import Buffer
-from zarr.core.buffer.core import default_buffer_prototype
+from zarr.core.buffer._core import default_buffer_prototype
 from zarr.core.common import concurrent_map
 
 if TYPE_CHECKING:

src/zarr/testing/store.py

@@ -10,7 +10,7 @@
     from typing import Any
 
     from zarr.abc.store import ByteRangeRequest
-    from zarr.core.buffer.core import BufferPrototype
+    from zarr.core.buffer._core import BufferPrototype
 
 import pytest
 

tests/test_codecs/test_codecs.py

@@ -23,7 +23,7 @@
 
 if TYPE_CHECKING:
     from zarr.abc.store import Store
-    from zarr.core.buffer.core import NDArrayLike
+    from zarr.core.buffer import NDArrayLike
     from zarr.core.common import MemoryOrder
 
 

tests/test_indexing.py

@@ -31,8 +31,7 @@
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator
 
-    from zarr.core.buffer import BufferPrototype
-    from zarr.core.buffer.core import Buffer
+    from zarr.core.buffer import Buffer, BufferPrototype
     from zarr.core.common import ChunkCoords