Skip to content

Commit f68bf06

Browse files
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
1 parent 7f4a681 commit f68bf06

File tree

14 files changed

+108
-24
lines changed

14 files changed

+108
-24
lines changed

changes/2871.feature.rst

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
Added public API for Buffer ABCs and implementations.
2+
3+
Use :mod:`zarr.buffer` to access buffer implementations, and
4+
:mod:`zarr.abc.buffer` for the interface to implement new buffer types.
5+
6+
Users previously importing buffer from ``zarr.core.buffer`` should update their
7+
imports to use :mod:`zarr.buffer`. As a reminder, all of ``zarr.core`` is
8+
considered a private API that's not covered by zarr-python's versioning policy.

docs/user-guide/config.rst

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ This is the current default configuration::
6363
'variable-length-string': {'name': 'vlen-utf8'}},
6464
'write_empty_chunks': False},
6565
'async': {'concurrency': 10, 'timeout': None},
66-
'buffer': 'zarr.core.buffer.cpu.Buffer',
66+
'buffer': 'zarr.buffer.cpu.Buffer',
6767
'codec_pipeline': {'batch_size': 1,
6868
'path': 'zarr.core.codec_pipeline.BatchedCodecPipeline'},
6969
'codecs': {'blosc': 'zarr.codecs.blosc.BloscCodec',
@@ -78,5 +78,5 @@ This is the current default configuration::
7878
'zstd': 'zarr.codecs.zstd.ZstdCodec'},
7979
'default_zarr_format': 3,
8080
'json_indent': 2,
81-
'ndbuffer': 'zarr.core.buffer.cpu.NDBuffer',
81+
'ndbuffer': 'zarr.buffer.cpu.NDBuffer',
8282
'threading': {'max_workers': None}}

docs/user-guide/extending.rst

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,10 @@ Coming soon.
8383
Custom array buffers
8484
--------------------
8585

86-
Coming soon.
86+
Zarr-python provides control over where and how arrays stored in memory through
87+
:mod:`zarr.buffer`. Currently both CPU (the default) and GPU implementations are
88+
provided (see :ref:`user-guide-gpu` for more). You can implement your own buffer
89+
classes by implementing the interface defined in :mod:`zarr.abc.buffer`.
8790

8891
Other extensions
8992
----------------

src/zarr/abc/buffer.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
from zarr.core.buffer.core import ArrayLike, Buffer, BufferPrototype, NDArrayLike, NDBuffer
2+
3+
__all__ = [
4+
"ArrayLike",
5+
"Buffer",
6+
"BufferPrototype",
7+
"NDArrayLike",
8+
"NDBuffer",
9+
]

src/zarr/buffer/__init__.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
"""
2+
Implementations of the Zarr Buffer interface.
3+
4+
See Also
5+
========
6+
zarr.abc.buffer: Abstract base class for the Zarr Buffer interface.
7+
"""
8+
9+
from zarr.buffer import cpu, gpu
10+
from zarr.core.buffer import default_buffer_prototype
11+
12+
__all__ = ["cpu", "default_buffer_prototype", "gpu"]

src/zarr/buffer/cpu.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
from zarr.core.buffer.cpu import (
2+
Buffer,
3+
NDBuffer,
4+
as_numpy_array_wrapper,
5+
buffer_prototype,
6+
numpy_buffer_prototype,
7+
)
8+
9+
__all__ = [
10+
"Buffer",
11+
"NDBuffer",
12+
"as_numpy_array_wrapper",
13+
"buffer_prototype",
14+
"numpy_buffer_prototype",
15+
]

src/zarr/buffer/gpu.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
from zarr.core.buffer.gpu import Buffer, NDBuffer, buffer_prototype
2+
3+
__all__ = [
4+
"Buffer",
5+
"NDBuffer",
6+
"buffer_prototype",
7+
]

src/zarr/core/buffer/cpu.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -224,5 +224,10 @@ def numpy_buffer_prototype() -> core.BufferPrototype:
224224
return core.BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
225225

226226

227-
register_buffer(Buffer)
228-
register_ndbuffer(NDBuffer)
227+
register_buffer(Buffer, qualname="zarr.buffer.cpu.Buffer")
228+
register_ndbuffer(NDBuffer, qualname="zarr.buffer.cpu.NDBuffer")
229+
230+
231+
# backwards compatibility
232+
register_buffer(Buffer, qualname="zarr.core.buffer.cpu.Buffer")
233+
register_ndbuffer(NDBuffer, qualname="zarr.core.buffer.cpu.NDBuffer")

src/zarr/core/buffer/gpu.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -220,5 +220,9 @@ def __setitem__(self, key: Any, value: Any) -> None:
220220

221221
buffer_prototype = BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
222222

223-
register_buffer(Buffer)
224-
register_ndbuffer(NDBuffer)
223+
register_buffer(Buffer, qualname="zarr.buffer.gpu.Buffer")
224+
register_ndbuffer(NDBuffer, qualname="zarr.buffer.gpu.NDBuffer")
225+
226+
# backwards compatibility
227+
register_buffer(Buffer, qualname="zarr.core.buffer.gpu.Buffer")
228+
register_ndbuffer(NDBuffer, qualname="zarr.core.buffer.gpu.NDBuffer")

src/zarr/core/config.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ def enable_gpu(self) -> ConfigSet:
7474
Configure Zarr to use GPUs where possible.
7575
"""
7676
return self.set(
77-
{"buffer": "zarr.core.buffer.gpu.Buffer", "ndbuffer": "zarr.core.buffer.gpu.NDBuffer"}
77+
{"buffer": "zarr.buffer.gpu.Buffer", "ndbuffer": "zarr.buffer.gpu.NDBuffer"}
7878
)
7979

8080

@@ -128,8 +128,8 @@ def enable_gpu(self) -> ConfigSet:
128128
"vlen-utf8": "zarr.codecs.vlen_utf8.VLenUTF8Codec",
129129
"vlen-bytes": "zarr.codecs.vlen_utf8.VLenBytesCodec",
130130
},
131-
"buffer": "zarr.core.buffer.cpu.Buffer",
132-
"ndbuffer": "zarr.core.buffer.cpu.NDBuffer",
131+
"buffer": "zarr.buffer.cpu.Buffer",
132+
"ndbuffer": "zarr.buffer.cpu.NDBuffer",
133133
}
134134
],
135135
)

0 commit comments

Comments
 (0)