add vlen-bytes

Author: d-v-b

src/zarr/core/dtype/__init__.py

@@ -12,6 +12,7 @@
     Structured,
 )
 from zarr.core.dtype.npy.time import DateTime64, TimeDelta64
+from zarr.core.dtype.npy.vlen_bytes import VariableLengthBytes
 
 if TYPE_CHECKING:
     from zarr.core.common import ZarrFormat
@@ -88,6 +89,7 @@
     | FixedLengthBytes
     | Structured
     | TimeDType
+    | VariableLengthBytes
 )
 # mypy has trouble inferring the type of variablelengthstring dtype, because its class definition
 # depends on the installed numpy version. That's why the type: ignore statement is needed here.
@@ -100,6 +102,7 @@
     FixedLengthBytes,
     Structured,
     *TIME_DTYPE,
+    VariableLengthBytes,
 )
 
 # This type models inputs that can be coerced to a ZDType

src/zarr/core/dtype/npy/vlen_bytes.py

@@ -1,15 +1,17 @@
+import base64
 from dataclasses import dataclass
 from typing import ClassVar, Literal, Self, TypeGuard, overload
 
 import numpy as np
 
 from zarr.core.common import JSON, ZarrFormat
-from zarr.core.dtype.common import HasObjectCodec
-from zarr.core.dtype.wrapper import TBaseDType, ZDType
+from zarr.core.dtype.common import HasObjectCodec, v3_unstable_dtype_warning
+from zarr.core.dtype.npy.common import check_json_str
+from zarr.core.dtype.wrapper import DTypeJSON_V2, DTypeJSON_V3, TBaseDType, ZDType
 
 
 @dataclass(frozen=True, kw_only=True)
-class VariableLengthString(ZDType[np.dtypes.ObjectDType, str], HasObjectCodec):  # type: ignore[no-redef]
+class VariableLengthBytes(ZDType[np.dtypes.ObjectDType, bytes], HasObjectCodec):
     dtype_cls = np.dtypes.ObjectDType
     _zarr_v3_name: ClassVar[Literal["variable_length_bytes"]] = "variable_length_bytes"
     object_codec_id = "vlen-bytes"
@@ -39,12 +41,13 @@ def check_json_v3(cls, data: JSON) -> TypeGuard[Literal["variable_length_utf8"]]
     def to_json(self, zarr_format: Literal[2]) -> Literal["|O"]: ...
 
     @overload
-    def to_json(self, zarr_format: Literal[3]) -> Literal["variable_length_utf8"]: ...
+    def to_json(self, zarr_format: Literal[3]) -> Literal["variable_length_bytes"]: ...
 
-    def to_json(self, zarr_format: ZarrFormat) -> Literal["|O", "variable_length_utf8"]:
+    def to_json(self, zarr_format: ZarrFormat) -> Literal["|O", "variable_length_bytes"]:
         if zarr_format == 2:
             return "|O"
         elif zarr_format == 3:
+            v3_unstable_dtype_warning(self)
             return self._zarr_v3_name
         raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}")  # pragma: no cover
 
@@ -54,22 +57,19 @@ def _from_json_unchecked(
     ) -> Self:
         return cls()
 
-    def default_scalar(self) -> str:
-        return ""
+    def default_scalar(self) -> bytes:
+        return b""
 
     def to_json_scalar(self, data: object, *, zarr_format: ZarrFormat) -> str:
-        return data  # type: ignore[return-value]
+        return base64.standard_b64encode(data).decode("ascii")  # type: ignore[arg-type]
 
-    def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> str:
-        """
-        Strings pass through
-        """
-        if not check_json_str(data):
-            raise TypeError(f"Invalid type: {data}. Expected a string.")
-        return data
+    def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> bytes:
+        if check_json_str(data):
+            return base64.standard_b64decode(data.encode("ascii"))
+        raise TypeError(f"Invalid type: {data}. Expected a string.")  # pragma: no cover
 
     def check_scalar(self, data: object) -> bool:
-        return isinstance(data, str)
+        return isinstance(data, bytes | str)
 
-    def _cast_scalar_unchecked(self, data: object) -> str:
-        return str(data)
+    def _cast_scalar_unchecked(self, data: object) -> bytes:
+        return bytes(data)  # type: ignore[no-any-return, call-overload]

src/zarr/core/dtype/wrapper.py

@@ -45,7 +45,7 @@
 
 # This the upper bound for the scalar types we support. It's numpy scalars + str,
 # because the new variable-length string dtype in numpy does not have a corresponding scalar type
-TBaseScalar = np.generic | str
+TBaseScalar = np.generic | str | bytes
 # This is the bound for the dtypes that we support. If we support non-numpy dtypes,
 # then this bound will need to be widened.
 TBaseDType = np.dtype[np.generic]
@@ -174,8 +174,8 @@ def cast_scalar(self, data: object) -> TScalar_co:
         if self.check_scalar(data):
             return self._cast_scalar_unchecked(data)
         msg = (
-            f"The value {data} failed a type check. "
-            f"It cannot be safely cast to a scalar compatible with {self.dtype_cls}. "
+            f"The value {data!r} failed a type check. "
+            f"It cannot be safely cast to a scalar compatible with {self}. "
             f"Consult the documentation for {self} to determine the possible values that can "
             "be cast to scalars of the wrapped data type."
         )

tests/test_regression/test_regression.py

@@ -13,6 +13,7 @@
 from zarr.core.array import Array
 from zarr.core.chunk_key_encodings import V2ChunkKeyEncoding
 from zarr.core.dtype.npy.string import VariableLengthString
+from zarr.core.dtype.npy.vlen_bytes import VariableLengthBytes
 from zarr.storage import LocalStore
 
 if TYPE_CHECKING:
@@ -33,7 +34,7 @@ def runner_installed() -> bool:
 @dataclass(kw_only=True)
 class ArrayParams:
     values: np.ndarray[tuple[int], np.dtype[np.generic]]
-    fill_value: np.generic | str | int
+    fill_value: np.generic | str | int | bytes
     filters: tuple[numcodecs.abc.Codec, ...] = ()
     compressor: numcodecs.abc.Codec
 
@@ -92,8 +93,10 @@ def source_array(tmp_path: Path, request: pytest.FixtureRequest) -> Array:
     compressor = array_params.compressor
     chunk_key_encoding = V2ChunkKeyEncoding(separator="/")
     dtype: ZDTypeLike
-    if array_params.values.dtype == np.dtype("|O"):
+    if array_params.values.dtype == np.dtype("|O") and array_params.filters == (VLenUTF8(),):
         dtype = VariableLengthString()  # type: ignore[assignment]
+    elif array_params.values.dtype == np.dtype("|O") and array_params.filters == (VLenBytes(),):
+        dtype = VariableLengthBytes()
     else:
         dtype = array_params.values.dtype
     z = zarr.create_array(