Public API for buffer objects (#2876)

* Public API for buffer objects

This moves the public imports from buffer things out of `zarr.core`.

Abstract stuff is availble under `zarr.abc.buffer`.

Concrete implementations are available under `zarr.buffer.{cpu,gpu}`.

* changelog

* absolute imports

* fixed warning in doc build

* Updated config

* Updated config

* wording

* doc

* Updated the test to reflect the break between class name and qualname

The config uses the public `zarr.buffer.cpu.Buffer`, which differs
from the implementation path `zarr.core.buffer.cpu.Buffer`. This is
OK because the public API for getting the buffer doesn't depend on
where it's implemented at.

* backwards compat

Author: TomAugspurger

changes/2871.feature.rst

@@ -0,0 +1,8 @@
+Added public API for Buffer ABCs and implementations.
+
+Use :mod:`zarr.buffer` to access buffer implementations, and
+:mod:`zarr.abc.buffer` for the interface to implement new buffer types.
+
+Users previously importing buffer from ``zarr.core.buffer`` should update their
+imports to use :mod:`zarr.buffer`. As a reminder, all of ``zarr.core`` is
+considered a private API that's not covered by zarr-python's versioning policy.

docs/user-guide/config.rst

@@ -63,7 +63,7 @@ This is the current default configuration::
                                        'variable-length-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',
@@ -78,5 +78,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}}

docs/user-guide/extending.rst

@@ -83,7 +83,10 @@ Coming soon.
 Custom array buffers
 --------------------
 
-Coming soon.
+Zarr-python provides control over 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 @@
+"""
+Implementations of the Zarr Buffer interface.
+
+See Also
+========
+zarr.abc.buffer: Abstract base class for the Zarr Buffer interface.
+"""
+
+from zarr.buffer import cpu, gpu
+from zarr.core.buffer import default_buffer_prototype
+
+__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",
+]

src/zarr/core/buffer/cpu.py

@@ -224,5 +224,10 @@ def numpy_buffer_prototype() -> core.BufferPrototype:
     return core.BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
 
 
-register_buffer(Buffer)
-register_ndbuffer(NDBuffer)
+register_buffer(Buffer, qualname="zarr.buffer.cpu.Buffer")
+register_ndbuffer(NDBuffer, qualname="zarr.buffer.cpu.NDBuffer")
+
+
+# backwards compatibility
+register_buffer(Buffer, qualname="zarr.core.buffer.cpu.Buffer")
+register_ndbuffer(NDBuffer, qualname="zarr.core.buffer.cpu.NDBuffer")

src/zarr/core/buffer/gpu.py

@@ -220,5 +220,9 @@ def __setitem__(self, key: Any, value: Any) -> None:
 
 buffer_prototype = BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
 
-register_buffer(Buffer)
-register_ndbuffer(NDBuffer)
+register_buffer(Buffer, qualname="zarr.buffer.gpu.Buffer")
+register_ndbuffer(NDBuffer, qualname="zarr.buffer.gpu.NDBuffer")
+
+# backwards compatibility
+register_buffer(Buffer, qualname="zarr.core.buffer.gpu.Buffer")
+register_ndbuffer(NDBuffer, qualname="zarr.core.buffer.gpu.NDBuffer")

src/zarr/core/config.py

@@ -74,7 +74,7 @@ def enable_gpu(self) -> ConfigSet:
         Configure Zarr to use GPUs where possible.
         """
         return self.set(
-            {"buffer": "zarr.core.buffer.gpu.Buffer", "ndbuffer": "zarr.core.buffer.gpu.NDBuffer"}
+            {"buffer": "zarr.buffer.gpu.Buffer", "ndbuffer": "zarr.buffer.gpu.NDBuffer"}
         )
 
 
@@ -128,8 +128,8 @@ def enable_gpu(self) -> ConfigSet:
                 "vlen-utf8": "zarr.codecs.vlen_utf8.VLenUTF8Codec",
                 "vlen-bytes": "zarr.codecs.vlen_utf8.VLenBytesCodec",
             },
-            "buffer": "zarr.core.buffer.cpu.Buffer",
-            "ndbuffer": "zarr.core.buffer.cpu.NDBuffer",
+            "buffer": "zarr.buffer.cpu.Buffer",
+            "ndbuffer": "zarr.buffer.cpu.NDBuffer",
         }
     ],
 )

src/zarr/core/group.py

@@ -1454,7 +1454,7 @@ async def create_hierarchy(
             group already exists at path ``a``, then this function will leave the group at ``a`` as-is.
 
         Yields
-        -------
+        ------
             tuple[str, AsyncArray | AsyncGroup].
         """
         # check that all the nodes have the same zarr_format as Self

src/zarr/registry.py

@@ -47,8 +47,10 @@ def lazy_load(self) -> None:
 
         self.lazy_load_list.clear()
 
-    def register(self, cls: type[T]) -> None:
-        self[fully_qualified_name(cls)] = cls
+    def register(self, cls: type[T], qualname: str | None = None) -> None:
+        if qualname is None:
+            qualname = fully_qualified_name(cls)
+        self[qualname] = cls
 
 
 __codec_registries: dict[str, Registry[Codec]] = defaultdict(Registry)
@@ -131,12 +133,12 @@ def register_pipeline(pipe_cls: type[CodecPipeline]) -> None:
     __pipeline_registry.register(pipe_cls)
 
 
-def register_ndbuffer(cls: type[NDBuffer]) -> None:
-    __ndbuffer_registry.register(cls)
+def register_ndbuffer(cls: type[NDBuffer], qualname: str | None = None) -> None:
+    __ndbuffer_registry.register(cls, qualname)
 
 
-def register_buffer(cls: type[Buffer]) -> None:
-    __buffer_registry.register(cls)
+def register_buffer(cls: type[Buffer], qualname: str | None = None) -> None:
+    __buffer_registry.register(cls, qualname)
 
 
 def get_codec_class(key: str, reload_config: bool = False) -> type[Codec]:

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,

tests/test_config.py

@@ -101,8 +101,8 @@ def test_config_defaults_set() -> None:
                     "vlen-utf8": "zarr.codecs.vlen_utf8.VLenUTF8Codec",
                     "vlen-bytes": "zarr.codecs.vlen_utf8.VLenBytesCodec",
                 },
-                "buffer": "zarr.core.buffer.cpu.Buffer",
-                "ndbuffer": "zarr.core.buffer.cpu.NDBuffer",
+                "buffer": "zarr.buffer.cpu.Buffer",
+                "ndbuffer": "zarr.buffer.cpu.NDBuffer",
             }
         ]
     )
@@ -224,9 +224,6 @@ class NewBloscCodec(BloscCodec):
 
 @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])
 def test_config_ndbuffer_implementation(store: Store) -> None:
-    # has default value
-    assert fully_qualified_name(get_ndbuffer_class()) == config.defaults[0]["ndbuffer"]
-
     # set custom ndbuffer with TestNDArrayLike implementation
     register_ndbuffer(NDBufferUsingTestNDArrayLike)
     with config.set({"ndbuffer": fully_qualified_name(NDBufferUsingTestNDArrayLike)}):
@@ -244,7 +241,7 @@ def test_config_ndbuffer_implementation(store: Store) -> None:
 
 def test_config_buffer_implementation() -> None:
     # has default value
-    assert fully_qualified_name(get_buffer_class()) == config.defaults[0]["buffer"]
+    assert config.defaults[0]["buffer"] == "zarr.buffer.cpu.Buffer"
 
     arr = zeros(shape=(100,), store=StoreExpectingTestBuffer())
 
@@ -279,6 +276,27 @@ def test_config_buffer_implementation() -> None:
         assert np.array_equal(arr_Crc32c[:], data2d)
 
 
+def test_config_buffer_backwards_compatibility() -> None:
+    # This should warn once zarr.core is private
+    # https://github.com/zarr-developers/zarr-python/issues/2621
+    with zarr.config.set(
+        {"buffer": "zarr.core.buffer.cpu.Buffer", "ndbuffer": "zarr.core.buffer.cpu.NDBuffer"}
+    ):
+        get_buffer_class()
+        get_ndbuffer_class()
+
+
+@pytest.mark.gpu
+def test_config_buffer_backwards_compatibility_gpu() -> None:
+    # This should warn once zarr.core is private
+    # https://github.com/zarr-developers/zarr-python/issues/2621
+    with zarr.config.set(
+        {"buffer": "zarr.core.buffer.gpu.Buffer", "ndbuffer": "zarr.core.buffer.gpu.NDBuffer"}
+    ):
+        get_buffer_class()
+        get_ndbuffer_class()
+
+
 @pytest.mark.filterwarnings("error")
 def test_warning_on_missing_codec_config() -> None:
     class NewCodec(BytesCodec):