From a9d2c9e44c068f8dcc72b23b1644a2ea8aff4c95 Mon Sep 17 00:00:00 2001
From: David Stansby <dstansby@gmail.com>
Date: Sat, 4 Jan 2025 09:39:35 +0000
Subject: [PATCH 1/2] Clean up zarr.core.buffer API

---
 src/zarr/abc/store.py                      |  2 +-
 src/zarr/core/buffer/__init__.py           |  2 +-
 src/zarr/core/buffer/{core.py => _core.py} |  3 ---
 src/zarr/core/buffer/cpu.py                | 24 ++++++++++++----------
 src/zarr/core/buffer/gpu.py                | 14 ++++++-------
 src/zarr/core/metadata/v3.py               |  2 +-
 src/zarr/storage/_local.py                 |  2 +-
 src/zarr/testing/store.py                  |  2 +-
 tests/test_codecs/test_codecs.py           |  2 +-
 tests/test_indexing.py                     |  3 +--
 tests/test_store/test_wrapper.py           |  2 +-
 11 files changed, 28 insertions(+), 30 deletions(-)
 rename src/zarr/core/buffer/{core.py => _core.py} (99%)

diff --git a/src/zarr/abc/store.py b/src/zarr/abc/store.py
index bd0a7ad503..2c82f1c752 100644
--- a/src/zarr/abc/store.py
+++ b/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
 
diff --git a/src/zarr/core/buffer/__init__.py b/src/zarr/core/buffer/__init__.py
index ccb41e291c..85486a2846 100644
--- a/src/zarr/core/buffer/__init__.py
+++ b/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,
diff --git a/src/zarr/core/buffer/core.py b/src/zarr/core/buffer/_core.py
similarity index 99%
rename from src/zarr/core/buffer/core.py
rename to src/zarr/core/buffer/_core.py
index 85a7351fc7..e9c8b8dcce 100644
--- a/src/zarr/core/buffer/core.py
+++ b/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):
diff --git a/src/zarr/core/buffer/cpu.py b/src/zarr/core/buffer/cpu.py
index 5019075496..1ed6b47795 100644
--- a/src/zarr/core/buffer/cpu.py
+++ b/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)
diff --git a/src/zarr/core/buffer/gpu.py b/src/zarr/core/buffer/gpu.py
index 6941c8897e..ed4c24ea9b 100644
--- a/src/zarr/core/buffer/gpu.py
+++ b/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)
diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py
index 0821dd9bc9..7cd7fe6783 100644
--- a/src/zarr/core/metadata/v3.py
+++ b/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
diff --git a/src/zarr/storage/_local.py b/src/zarr/storage/_local.py
index f4226792cb..15842ce959 100644
--- a/src/zarr/storage/_local.py
+++ b/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:
diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py
index ada028c273..bbf785aa83 100644
--- a/src/zarr/testing/store.py
+++ b/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
 
diff --git a/tests/test_codecs/test_codecs.py b/tests/test_codecs/test_codecs.py
index e36a332440..b8122b4ac2 100644
--- a/tests/test_codecs/test_codecs.py
+++ b/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
 
 
diff --git a/tests/test_indexing.py b/tests/test_indexing.py
index 30d0d75f22..1317f13e46 100644
--- a/tests/test_indexing.py
+++ b/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
 
 
diff --git a/tests/test_store/test_wrapper.py b/tests/test_store/test_wrapper.py
index 489bcd5a7a..6d548568fb 100644
--- a/tests/test_store/test_wrapper.py
+++ b/tests/test_store/test_wrapper.py
@@ -9,7 +9,7 @@
 
 if TYPE_CHECKING:
     from zarr.abc.store import Store
-    from zarr.core.buffer.core import BufferPrototype
+    from zarr.core.buffer import BufferPrototype
 
 
 @pytest.mark.parametrize("store", ["local", "memory", "zip"], indirect=True)

From 52d51bc231c9a64d28f2baa5a0c399f1e9439238 Mon Sep 17 00:00:00 2001
From: David Stansby <dstansby@gmail.com>
Date: Sat, 4 Jan 2025 09:56:30 +0000
Subject: [PATCH 2/2] Fix circular import

---
 src/zarr/codecs/sharding.py      | 2 +-
 src/zarr/core/buffer/__init__.py | 2 --
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py
index a01145b3b2..299ea65df9 100644
--- a/src/zarr/codecs/sharding.py
+++ b/src/zarr/codecs/sharding.py
@@ -26,8 +26,8 @@
     BufferPrototype,
     NDBuffer,
     default_buffer_prototype,
-    numpy_buffer_prototype,
 )
+from zarr.core.buffer.cpu import numpy_buffer_prototype
 from zarr.core.chunk_grids import ChunkGrid, RegularChunkGrid
 from zarr.core.common import (
     ChunkCoords,
diff --git a/src/zarr/core/buffer/__init__.py b/src/zarr/core/buffer/__init__.py
index 85486a2846..d3df13129f 100644
--- a/src/zarr/core/buffer/__init__.py
+++ b/src/zarr/core/buffer/__init__.py
@@ -6,7 +6,6 @@
     NDBuffer,
     default_buffer_prototype,
 )
-from zarr.core.buffer.cpu import numpy_buffer_prototype
 
 __all__ = [
     "ArrayLike",
@@ -15,5 +14,4 @@
     "NDArrayLike",
     "NDBuffer",
     "default_buffer_prototype",
-    "numpy_buffer_prototype",
 ]