Public API for buffer objects

changes/2871.feature.rst

@@ -0,0 +1 @@
+Add public :mod:`zarr.buffer` API for controlling how data is stored in memory.

docs/user-guide/extending.rst

@@ -83,7 +83,10 @@ Coming soon.
 Custom array buffers
 --------------------
 
-Coming soon.
+zarr-python provides control where and how arrays stored in memory through
+:mod:`zarr.buffer`. Currently both CPU (the default) and GPU implementations are
+provided (see :ref:`user-guide-gpu` for more). You can implement your own buffer
+classes by implementing the interface defined in :mod:`zarr.abc.buffer`.
 
 Other extensions
 ----------------

src/zarr/abc/buffer.py

@@ -0,0 +1,9 @@
+from zarr.core.buffer.core import ArrayLike, Buffer, BufferPrototype, NDArrayLike, NDBuffer
+
+__all__ = [
+    "ArrayLike",
+    "Buffer",
+    "BufferPrototype",
+    "NDArrayLike",
+    "NDBuffer",
+]

src/zarr/buffer/__init__.py

@@ -0,0 +1,12 @@
+"""
+Public API for implementations of the Zarr Buffer interface.
+
+See Also
+========
+arr.abc.buffer: Abstract base class for the Zarr Buffer interface.
+"""
+
+from ..core.buffer import default_buffer_prototype
+from . import cpu, gpu
+
+__all__ = ["cpu", "default_buffer_prototype", "gpu"]

src/zarr/buffer/cpu.py

@@ -0,0 +1,15 @@
+from zarr.core.buffer.cpu import (
+    Buffer,
+    NDBuffer,
+    as_numpy_array_wrapper,
+    buffer_prototype,
+    numpy_buffer_prototype,
+)
+
+__all__ = [
+    "Buffer",
+    "NDBuffer",
+    "as_numpy_array_wrapper",
+    "buffer_prototype",
+    "numpy_buffer_prototype",
+]

src/zarr/buffer/gpu.py

@@ -0,0 +1,7 @@
+from zarr.core.buffer.gpu import Buffer, NDBuffer, buffer_prototype
+
+__all__ = [
+    "Buffer",
+    "NDBuffer",
+    "buffer_prototype",
+]

tests/test_buffer.py

@@ -6,12 +6,13 @@
 import pytest
 
 import zarr
+from zarr.abc.buffer import ArrayLike, BufferPrototype, NDArrayLike
+from zarr.buffer import cpu, gpu
 from zarr.codecs.blosc import BloscCodec
 from zarr.codecs.crc32c_ import Crc32cCodec
 from zarr.codecs.gzip import GzipCodec
 from zarr.codecs.transpose import TransposeCodec
 from zarr.codecs.zstd import ZstdCodec
-from zarr.core.buffer import ArrayLike, BufferPrototype, NDArrayLike, cpu, gpu
 from zarr.storage import MemoryStore, StorePath
 from zarr.testing.buffer import (
     NDBufferUsingTestNDArrayLike,