From 98624241e9ab435f2e8fd365be2af828c1a66ce5 Mon Sep 17 00:00:00 2001 From: Chris Barnes Date: Wed, 21 Oct 2026 10:37:46 +0100 Subject: [PATCH 1/4] Implement and test the reshape codec Based on external implementation at https://github.com/clbarnes/zarr-python-reshape . --- src/zarr/codecs/__init__.py | 3 + src/zarr/codecs/reshape.py | 272 ++++++++++++++++++++++++++++++ tests/test_codecs/test_reshape.py | 94 +++++++++++ 3 files changed, 369 insertions(+) create mode 100644 src/zarr/codecs/reshape.py create mode 100644 tests/test_codecs/test_reshape.py diff --git a/src/zarr/codecs/__init__.py b/src/zarr/codecs/__init__.py index 9a1b47b351..ecd06ee0f5 100644 --- a/src/zarr/codecs/__init__.py +++ b/src/zarr/codecs/__init__.py @@ -28,6 +28,7 @@ Zlib, Zstd, ) +from zarr.codecs.reshape import ReshapeCodec from zarr.codecs.scale_offset import ScaleOffset from zarr.codecs.sharding import ShardingCodec, ShardingCodecIndexLocation, SubchunkWriteOrder from zarr.codecs.transpose import TransposeCodec @@ -44,6 +45,7 @@ "Crc32cCodec", "Endian", "GzipCodec", + "ReshapeCodec", "ScaleOffset", "ShardingCodec", "ShardingCodecIndexLocation", @@ -57,6 +59,7 @@ register_codec("blosc", BloscCodec) register_codec("cast_value", CastValue) register_codec("bytes", BytesCodec) +register_codec("reshape", ReshapeCodec) # compatibility with earlier versions of ZEP1 register_codec("endian", BytesCodec) diff --git a/src/zarr/codecs/reshape.py b/src/zarr/codecs/reshape.py new file mode 100644 index 0000000000..84b69a6c92 --- /dev/null +++ b/src/zarr/codecs/reshape.py @@ -0,0 +1,272 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from zarr.abc.codec import ArrayArrayCodec +from zarr.core.array_spec import ArraySpec +from zarr.core.common import JSON, parse_named_configuration + +if TYPE_CHECKING: + from collections.abc import Iterable, Sequence + from typing import Self + + from zarr.core.buffer import NDBuffer + from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType + from zarr.core.metadata.v3 import ChunkGridMetadata + + +# Type for shape specification: int, list of ints, or -1 +type ShapeElement = int | tuple[int, ...] + + +def prod(values: Iterable[int]) -> int: + p = 1 + for v in values: + p *= v + return p + + +def _compute_output_shape( + shape_spec: Sequence[ShapeElement], + input_shape: Sequence[int], +) -> tuple[int, ...]: + """ + Compute the output shape from the shape specification and input shape. + + Parameters + ---------- + shape_spec : tuple[ShapeElement, ...] + The shape specification from the codec configuration. + input_shape : tuple[int, ...] + The shape of the input array. + + Returns + ------- + tuple[int, ...] + The computed output shape. + + Raises + ------ + ValueError + If the shape specification is invalid or the invariant + prod(output_shape) == prod(input_shape) cannot be satisfied. + """ + _validate_shape_spec(shape_spec, len(input_shape)) + output_shape: list[int] = [] + infer_idx: int | None = None + input_total = int(prod(input_shape)) + + # we can skip a lot of validity checks here because we have already used _validate_shape_spec + for i, elem in enumerate(shape_spec): + if isinstance(elem, int): + if elem == -1: + if infer_idx is not None: + raise ValueError( + "The special value -1 may occur at most once in the shape specification." + ) + infer_idx = i + output_shape.append(-1) # placeholder + else: + output_shape.append(elem) + else: + # Product of input dimensions specified + output_shape.append(prod(input_shape[d] for d in elem)) + + # Handle inference of -1 dimension + if infer_idx is not None: + known_product = prod(s for s in output_shape if s != -1) + if known_product == 0 or input_total % known_product != 0: + raise ValueError( + f"Cannot infer dimension size: input total {input_total} is not divisible " + f"by known product {known_product}." + ) + output_shape[infer_idx] = input_total // known_product + + output_total = prod(output_shape) + if output_total != input_total: + raise ValueError( + f"The invariant prod(output_shape) == prod(input_shape) is not satisfied. " + f"Got prod({output_shape}) = {output_total} != prod({input_shape}) = {input_total}." + ) + + return tuple(output_shape) + + +def _validate_shape_spec(shape_spec: Sequence[ShapeElement], input_ndim: int | None = None) -> None: + """ + Validate the shape specification according to the constraints. + + Validates: + 1. Input dimensions are in monotonically increasing order. + 2. For each input_dims list with k > 0 dimensions, the constraints about + raveled indices are satisfiable. + + Parameters + ---------- + shape_spec : tuple[ShapeElement, ...] + The shape specification from the codec configuration. + input_ndim : int or None + Number of dimensions in the input array, if available. + + Raises + ------ + ValueError + If any constraint is violated. + """ + if input_ndim is None: + dim_idxs = None + else: + dim_idxs = set(range(input_ndim)) + + last_dim_idx = -1 + found_minus_one = False + msg = "Arrays in reshape shape spec must contain unique indices into the input shape, in strictly monotonic increasing order" + + for elem in shape_spec: + if isinstance(elem, int): + pass + else: + for d in elem: + if d <= last_dim_idx: + raise ValueError(msg) + + last_dim_idx = d + + if d == -1: + if found_minus_one: + raise ValueError("-1 should appear at most once in reshape spec") + found_minus_one = True + elif d < 0: + raise ValueError("Negative indices not allowed in reshape spec") + + if dim_idxs is not None: + try: + dim_idxs.remove(d) + except KeyError as e: + raise ValueError(msg) from e + + +@dataclass(frozen=True) +class ReshapeCodec(ArrayArrayCodec): + """Reshape codec. + + An array-to-array codec that performs a reshape operation, similar to numpy's + reshape function. The reshape operation preserves the lexicographical (C-order) + traversal of elements. + + Parameters + ---------- + shape : Iterable[int | list[int]] + An array specifying the size of each dimension of the output array as a + function of the input array shape. Each element must be one of: + - A positive integer specifying an explicit size. + - A list of integers specifying input dimensions whose product determines + the output dimension size. + - The special value -1 (at most once), which is inferred to satisfy + prod(output_shape) == prod(input_shape). + + Notes + ----- + This codec does NOT alter the lexicographical order of elements. The contents + of the output array B is related to the input array A by: + ``ravel(B) == ravel(A)``. + + When possible, implementations should construct a virtual view rather than + copy the array. + + References + ---------- + + - The `reshape` codec spec: https://github.com/zarr-developers/zarr-extensions/tree/main/codecs/scale_offset + """ + + is_fixed_size = True + + shape: tuple[ShapeElement, ...] + + def __init__(self, *, shape: tuple[ShapeElement, ...]) -> None: + _validate_shape_spec(shape) + object.__setattr__(self, "shape", shape) + + @classmethod + def from_dict(cls, data: dict[str, JSON]) -> Self: + _, configuration_parsed = parse_named_configuration(data, "reshape") + return cls(**configuration_parsed) # type: ignore[arg-type] + + def to_dict(self) -> dict[str, JSON]: + return {"name": "reshape", "configuration": {"shape": self.shape}} + + def validate( + self, + shape: tuple[int, ...], + dtype: ZDType[TBaseDType, TBaseScalar], + chunk_grid: ChunkGridMetadata, + ) -> None: + """Validate the codec configuration against the array metadata.""" + _compute_output_shape(self.shape, shape) + + def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self: + """Evolve the codec from an array specification.""" + _compute_output_shape(self.shape, array_spec.shape) + return self + + def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: + """Resolve the output metadata after applying this codec.""" + output_shape = _compute_output_shape(self.shape, chunk_spec.shape) + return ArraySpec( + shape=output_shape, + dtype=chunk_spec.dtype, + fill_value=chunk_spec.fill_value, + config=chunk_spec.config, + prototype=chunk_spec.prototype, + ) + + def _decode_sync( + self, + chunk_array: NDBuffer, + chunk_spec: ArraySpec, + ) -> NDBuffer: + """Decode: reshape back to the original shape.""" + # chunk_spec contains the *decoded* (original) shape + # The input chunk_array has the encoded (reshaped) shape + # We need to reshape it back to the original shape + inferred_shape = _compute_output_shape(self.shape, chunk_spec.shape) + if inferred_shape != chunk_array.shape: + raise RuntimeError( + "Shape of chunk to decode %s does not match shape of requested array %s encoded with reshape codec shape=%s", + chunk_array.shape, + chunk_spec.shape, + self.shape, + ) + return chunk_array.reshape(chunk_spec.shape) + + async def _decode_single( + self, + chunk_array: NDBuffer, + chunk_spec: ArraySpec, + ) -> NDBuffer: + return self._decode_sync(chunk_array, chunk_spec) + + def _encode_sync( + self, + chunk_array: NDBuffer, + chunk_spec: ArraySpec, + ) -> NDBuffer | None: + """Encode: reshape to the output shape.""" + output_shape = _compute_output_shape(self.shape, chunk_spec.shape) + return chunk_array.reshape(output_shape) + + async def _encode_single( + self, + chunk_array: NDBuffer, + chunk_spec: ArraySpec, + ) -> NDBuffer | None: + return self._encode_sync(chunk_array, chunk_spec) + + def compute_encoded_size(self, input_byte_length: int, _chunk_spec: ArraySpec) -> int: + """Compute the encoded size (same as input since reshape doesn't change data).""" + return input_byte_length + + +__all__ = ["ReshapeCodec"] diff --git a/tests/test_codecs/test_reshape.py b/tests/test_codecs/test_reshape.py new file mode 100644 index 0000000000..6a8a57dd38 --- /dev/null +++ b/tests/test_codecs/test_reshape.py @@ -0,0 +1,94 @@ +import json + +import numpy as np +import pytest + +from zarr.codecs.reshape import ReshapeCodec +from zarr.core.array_spec import ArrayConfig, ArraySpec +from zarr.core.buffer import default_buffer_prototype +from zarr.core.dtype import UInt8 + + +def test_deserialize() -> None: + s = '{"name": "reshape", "configuration": {"shape": [5, [1, 2], -1]}}' + d = json.loads(s) + ReshapeCodec.from_dict(d) + + +@pytest.mark.asyncio +async def test_encode() -> None: + arr_flat = np.arange(60, dtype="uint8") + arr_3x2 = np.reshape(arr_flat, (3, 4, 5)) + dt = UInt8() + codec = ReshapeCodec(shape=((0, 1), 5)) + proto = default_buffer_prototype() + (encoded, *_) = await codec.encode( + [ + ( + proto.nd_buffer.from_numpy_array(arr_3x2), + ArraySpec( + arr_3x2.shape, + dt, + fill_value=0, + config=ArrayConfig("C", False), + prototype=proto, + ), + ) + ] + ) + assert encoded is not None + assert encoded.shape == (12, 5) + assert np.all(np.equal(np.ravel(encoded.as_numpy_array()), arr_flat)) + + +@pytest.mark.asyncio +async def test_decode() -> None: + arr_flat = np.arange(60, dtype="uint8") + arr_12x5 = np.reshape(arr_flat, (12, 5)) + dt = UInt8() + codec = ReshapeCodec(shape=((0, 1), 5)) + proto = default_buffer_prototype() + (encoded, *_) = await codec.decode( + [ + ( + proto.nd_buffer.from_numpy_array(arr_12x5), + ArraySpec( + (3, 4, 5), + dt, + fill_value=0, + config=ArrayConfig("C", False), + prototype=proto, + ), + ) + ] + ) + assert encoded is not None + assert encoded.shape == (3, 4, 5) + assert np.all(np.equal(np.ravel(encoded.as_numpy_array()), arr_flat)) + + +@pytest.mark.asyncio +async def test_example() -> None: + import json + + import zarr + from zarr.storage import ManagedMemoryStore + + # Create an array with the reshape codec + # This reshapes chunks from (100, 50, 64, 3) to (5000, 64, 3) + store = ManagedMemoryStore() + _ = zarr.create_array( + store, + shape=(1000, 500, 64, 3), + chunks=(100, 50, 64, 3), + dtype="float32", + filters=[ReshapeCodec(shape=((0, 1), (2,), 3))], + ) + b = await store.get("zarr.json") + assert b is not None + s = b.to_bytes().decode("utf-8") + d = json.loads(s) + assert d["codecs"][0] == { + "name": "reshape", + "configuration": {"shape": [[0, 1], [2], 3]}, + } From ab39719b87f0c226b1a6284b78f7f52f39be7854 Mon Sep 17 00:00:00 2001 From: Chris Barnes Date: Wed, 21 Oct 2026 10:44:03 +0100 Subject: [PATCH 2/4] Reshape changelog --- changes/4542.feature.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/4542.feature.md diff --git a/changes/4542.feature.md b/changes/4542.feature.md new file mode 100644 index 0000000000..8cba486083 --- /dev/null +++ b/changes/4542.feature.md @@ -0,0 +1 @@ +Implement the [reshape extension codec](https://github.com/zarr-developers/zarr-extensions/tree/main/codecs/reshape). From cde8c5d508822673282a26fba366c17189342ab6 Mon Sep 17 00:00:00 2001 From: Chris Barnes Date: Wed, 21 Oct 2026 13:57:14 +0100 Subject: [PATCH 3/4] Fix numpydoc --- src/zarr/codecs/reshape.py | 267 ++++++++++++++---------------- tests/test_codecs/test_reshape.py | 83 +++++++++- 2 files changed, 208 insertions(+), 142 deletions(-) diff --git a/src/zarr/codecs/reshape.py b/src/zarr/codecs/reshape.py index 84b69a6c92..b52ab8a688 100644 --- a/src/zarr/codecs/reshape.py +++ b/src/zarr/codecs/reshape.py @@ -18,6 +18,7 @@ # Type for shape specification: int, list of ints, or -1 type ShapeElement = int | tuple[int, ...] +type ShapeSpec = tuple[ShapeElement, ...] def prod(values: Iterable[int]) -> int: @@ -27,126 +28,6 @@ def prod(values: Iterable[int]) -> int: return p -def _compute_output_shape( - shape_spec: Sequence[ShapeElement], - input_shape: Sequence[int], -) -> tuple[int, ...]: - """ - Compute the output shape from the shape specification and input shape. - - Parameters - ---------- - shape_spec : tuple[ShapeElement, ...] - The shape specification from the codec configuration. - input_shape : tuple[int, ...] - The shape of the input array. - - Returns - ------- - tuple[int, ...] - The computed output shape. - - Raises - ------ - ValueError - If the shape specification is invalid or the invariant - prod(output_shape) == prod(input_shape) cannot be satisfied. - """ - _validate_shape_spec(shape_spec, len(input_shape)) - output_shape: list[int] = [] - infer_idx: int | None = None - input_total = int(prod(input_shape)) - - # we can skip a lot of validity checks here because we have already used _validate_shape_spec - for i, elem in enumerate(shape_spec): - if isinstance(elem, int): - if elem == -1: - if infer_idx is not None: - raise ValueError( - "The special value -1 may occur at most once in the shape specification." - ) - infer_idx = i - output_shape.append(-1) # placeholder - else: - output_shape.append(elem) - else: - # Product of input dimensions specified - output_shape.append(prod(input_shape[d] for d in elem)) - - # Handle inference of -1 dimension - if infer_idx is not None: - known_product = prod(s for s in output_shape if s != -1) - if known_product == 0 or input_total % known_product != 0: - raise ValueError( - f"Cannot infer dimension size: input total {input_total} is not divisible " - f"by known product {known_product}." - ) - output_shape[infer_idx] = input_total // known_product - - output_total = prod(output_shape) - if output_total != input_total: - raise ValueError( - f"The invariant prod(output_shape) == prod(input_shape) is not satisfied. " - f"Got prod({output_shape}) = {output_total} != prod({input_shape}) = {input_total}." - ) - - return tuple(output_shape) - - -def _validate_shape_spec(shape_spec: Sequence[ShapeElement], input_ndim: int | None = None) -> None: - """ - Validate the shape specification according to the constraints. - - Validates: - 1. Input dimensions are in monotonically increasing order. - 2. For each input_dims list with k > 0 dimensions, the constraints about - raveled indices are satisfiable. - - Parameters - ---------- - shape_spec : tuple[ShapeElement, ...] - The shape specification from the codec configuration. - input_ndim : int or None - Number of dimensions in the input array, if available. - - Raises - ------ - ValueError - If any constraint is violated. - """ - if input_ndim is None: - dim_idxs = None - else: - dim_idxs = set(range(input_ndim)) - - last_dim_idx = -1 - found_minus_one = False - msg = "Arrays in reshape shape spec must contain unique indices into the input shape, in strictly monotonic increasing order" - - for elem in shape_spec: - if isinstance(elem, int): - pass - else: - for d in elem: - if d <= last_dim_idx: - raise ValueError(msg) - - last_dim_idx = d - - if d == -1: - if found_minus_one: - raise ValueError("-1 should appear at most once in reshape spec") - found_minus_one = True - elif d < 0: - raise ValueError("Negative indices not allowed in reshape spec") - - if dim_idxs is not None: - try: - dim_idxs.remove(d) - except KeyError as e: - raise ValueError(msg) from e - - @dataclass(frozen=True) class ReshapeCodec(ArrayArrayCodec): """Reshape codec. @@ -155,16 +36,12 @@ class ReshapeCodec(ArrayArrayCodec): reshape function. The reshape operation preserves the lexicographical (C-order) traversal of elements. - Parameters - ---------- - shape : Iterable[int | list[int]] - An array specifying the size of each dimension of the output array as a - function of the input array shape. Each element must be one of: - - A positive integer specifying an explicit size. - - A list of integers specifying input dimensions whose product determines - the output dimension size. - - The special value -1 (at most once), which is inferred to satisfy - prod(output_shape) == prod(input_shape). + The given `shape`'s elements must be one of: + - A positive integer specifying an explicit size. + - A list of integers specifying input dimensions whose product determines + the output dimension size. + - The special value -1 (at most once), which is inferred to satisfy + prod(output_shape) == prod(input_shape). Notes ----- @@ -178,16 +55,15 @@ class ReshapeCodec(ArrayArrayCodec): References ---------- - - The `reshape` codec spec: https://github.com/zarr-developers/zarr-extensions/tree/main/codecs/scale_offset + - The `reshape` codec spec: https://github.com/zarr-developers/zarr-extensions/tree/main/codecs/reshape """ is_fixed_size = True - shape: tuple[ShapeElement, ...] + shape: ShapeSpec - def __init__(self, *, shape: tuple[ShapeElement, ...]) -> None: - _validate_shape_spec(shape) - object.__setattr__(self, "shape", shape) + def __post_init__(self) -> None: + self._validate_shape_spec() @classmethod def from_dict(cls, data: dict[str, JSON]) -> Self: @@ -204,16 +80,16 @@ def validate( chunk_grid: ChunkGridMetadata, ) -> None: """Validate the codec configuration against the array metadata.""" - _compute_output_shape(self.shape, shape) + self._compute_output_shape(shape) def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self: """Evolve the codec from an array specification.""" - _compute_output_shape(self.shape, array_spec.shape) + self._compute_output_shape(array_spec.shape) return self def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: """Resolve the output metadata after applying this codec.""" - output_shape = _compute_output_shape(self.shape, chunk_spec.shape) + output_shape = self._compute_output_shape(chunk_spec.shape) return ArraySpec( shape=output_shape, dtype=chunk_spec.dtype, @@ -231,7 +107,7 @@ def _decode_sync( # chunk_spec contains the *decoded* (original) shape # The input chunk_array has the encoded (reshaped) shape # We need to reshape it back to the original shape - inferred_shape = _compute_output_shape(self.shape, chunk_spec.shape) + inferred_shape = self._compute_output_shape(chunk_spec.shape) if inferred_shape != chunk_array.shape: raise RuntimeError( "Shape of chunk to decode %s does not match shape of requested array %s encoded with reshape codec shape=%s", @@ -254,7 +130,7 @@ def _encode_sync( chunk_spec: ArraySpec, ) -> NDBuffer | None: """Encode: reshape to the output shape.""" - output_shape = _compute_output_shape(self.shape, chunk_spec.shape) + output_shape = self._compute_output_shape(chunk_spec.shape) return chunk_array.reshape(output_shape) async def _encode_single( @@ -264,9 +140,118 @@ async def _encode_single( ) -> NDBuffer | None: return self._encode_sync(chunk_array, chunk_spec) - def compute_encoded_size(self, input_byte_length: int, _chunk_spec: ArraySpec) -> int: + def compute_encoded_size(self, input_byte_length: int, chunk_spec: ArraySpec) -> int: """Compute the encoded size (same as input since reshape doesn't change data).""" return input_byte_length + def _compute_output_shape( + self, + input_shape: Sequence[int], + ) -> tuple[int, ...]: + """ + Compute the output shape from the shape specification and input shape. + + Parameters + ---------- + input_shape : tuple[int, ...] + The shape of the input array. + + Returns + ------- + tuple[int, ...] + The computed output shape. + + Raises + ------ + ValueError + If the shape specification is invalid or the invariant + prod(output_shape) == prod(input_shape) cannot be satisfied. + """ + output_shape: list[int] = [] + infer_idx: int | None = None + input_total = int(prod(input_shape)) + + # we can skip a lot of validity checks here because we have already used _validate_shape_spec + for i, elem in enumerate(self.shape): + if isinstance(elem, int): + if elem == -1: + infer_idx = i + output_shape.append(-1) # placeholder + else: + output_shape.append(elem) + continue + + os = 1 + + for d in elem: + if d >= len(input_shape): + raise ValueError(f"Dimension index {d} out of range for reshape spec") + os *= input_shape[d] + + # Product of input dimensions specified + output_shape.append(os) + + # Handle inference of -1 dimension + if infer_idx is not None: + known_product = prod(s for s in output_shape if s != -1) + if known_product == 0 or input_total % known_product != 0: + raise ValueError( + f"Cannot infer dimension size: input total {input_total} is not divisible " + f"by known product {known_product}." + ) + output_shape[infer_idx] = input_total // known_product + + output_total = prod(output_shape) + if output_total != input_total: + raise ValueError( + f"The invariant prod(output_shape) == prod(input_shape) is not satisfied. " + f"Got prod({output_shape}) = {output_total} != prod({input_shape}) = {input_total}." + ) + + return tuple(output_shape) + + def _validate_shape_spec(self) -> None: + """ + Validate the shape specification according to the constraints. + + Validates: + 1. Input dimensions are in monotonically increasing order. + 2. For each input_dims list with k > 0 dimensions, the constraints about + raveled indices are satisfiable. + + Raises + ------ + ValueError + If more than explicit length of -1 is found, + or any explicit length < -1 is found, + or any indices into the array shape are below 0, + or the indices into the array shape are not monotonically increasing. + """ + last_dim_idx = -1 + found_minus_one = False + + for elem in self.shape: + if isinstance(elem, int): + if elem == -1: + if found_minus_one: + raise ValueError("-1 should appear at most once in the reshape spec") + else: + found_minus_one = True + elif elem < 1: + raise ValueError("Explicit axis size must be >=1 in the reshape spec") + continue + + for d in elem: + if d < 0: + raise ValueError( + "Index of input dimension must not be negative in the reshape spec" + ) + if d <= last_dim_idx: + raise ValueError( + "Index of input dimension must be monotonically increasing in reshape spec" + ) + + last_dim_idx = d + __all__ = ["ReshapeCodec"] diff --git a/tests/test_codecs/test_reshape.py b/tests/test_codecs/test_reshape.py index 6a8a57dd38..a8a4fa111f 100644 --- a/tests/test_codecs/test_reshape.py +++ b/tests/test_codecs/test_reshape.py @@ -3,10 +3,11 @@ import numpy as np import pytest -from zarr.codecs.reshape import ReshapeCodec +from zarr.codecs.reshape import ReshapeCodec, ShapeSpec from zarr.core.array_spec import ArrayConfig, ArraySpec from zarr.core.buffer import default_buffer_prototype from zarr.core.dtype import UInt8 +from zarr.core.metadata.v3 import RegularChunkGridMetadata def test_deserialize() -> None: @@ -67,6 +68,86 @@ async def test_decode() -> None: assert np.all(np.equal(np.ravel(encoded.as_numpy_array()), arr_flat)) +@pytest.mark.parametrize( + "shape_spec", + [ + (-2,), # negative number isn't -1 + (-1, -1), # more than one -1 + ((-2, 1),), # negative index + ((1, 0),), # decreasing + # (((0, 4),), 3), # missing dim + ], +) +def test_invalid_shape_spec(shape_spec: ShapeSpec) -> None: + with pytest.raises(ValueError): + ReshapeCodec(shape=shape_spec) + + +def test_compute_encoded_size() -> None: + c = ReshapeCodec(shape=((0, 1), 5)) + + dec_shape = (64, 64) + dec_size = np.prod(dec_shape).item() + enc_size = c.compute_encoded_size( + dec_size, + # ignored + ArraySpec( + (64, 64), + UInt8(), + fill_value=0, + config=ArrayConfig("C", False), + prototype=default_buffer_prototype(), + ), + ) + assert dec_size == enc_size + + +@pytest.mark.parametrize( + ("shape", "shape_spec"), + [ + ((3, 4), ((0, 1, 2),)), # index out of range + ((3, 4), (5, -1)), # indivisible -1 + ((3, 4), (5, 2)), # wrong shape + ], +) +def test_validate_invalid(shape: tuple[int, ...], shape_spec: ShapeSpec) -> None: + """Stand-in for _compute_output_shape""" + codec = ReshapeCodec(shape=shape_spec) + with pytest.raises(ValueError): + codec.validate( + shape, + UInt8(), # ignored + RegularChunkGridMetadata(chunk_shape=shape), + ) + + +@pytest.mark.parametrize( + ("input_shape", "shape_spec", "output_shape"), + [ + ((3, 4), (4, 3), (4, 3)), + ((3, 4), (6, -1), (6, 2)), + ((3, 4), ((0, 1),), (12,)), + ], +) +def test_valid( + input_shape: tuple[int, ...], shape_spec: ShapeSpec, output_shape: tuple[int, ...] +) -> None: + codec = ReshapeCodec(shape=shape_spec) + dtype = UInt8() + codec.validate(input_shape, dtype, RegularChunkGridMetadata(chunk_shape=input_shape)) + input_spec = ArraySpec( + input_shape, dtype, 0, ArrayConfig("C", False), default_buffer_prototype() + ) + output_spec = codec.resolve_metadata(input_spec) + + assert output_spec.dtype == input_spec.dtype + assert output_spec.config == input_spec.config + assert output_spec.fill_value == input_spec.fill_value + assert output_spec.prototype == input_spec.prototype + + assert output_spec.shape == output_shape + + @pytest.mark.asyncio async def test_example() -> None: import json From 40e89f7cdf49c88990757bf348e5c13ab15535cb Mon Sep 17 00:00:00 2001 From: Chris Barnes Date: Thu, 22 Oct 2026 09:40:59 +0100 Subject: [PATCH 4/4] Fix reshape test coverage --- tests/test_codecs/test_reshape.py | 33 ++++++++++++++++--------------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/tests/test_codecs/test_reshape.py b/tests/test_codecs/test_reshape.py index a8a4fa111f..2803c34c59 100644 --- a/tests/test_codecs/test_reshape.py +++ b/tests/test_codecs/test_reshape.py @@ -10,6 +10,10 @@ from zarr.core.metadata.v3 import RegularChunkGridMetadata +def array_spec(shape: tuple[int, ...]) -> ArraySpec: + return ArraySpec(shape, UInt8(), 0, ArrayConfig("C", False), default_buffer_prototype()) + + def test_deserialize() -> None: s = '{"name": "reshape", "configuration": {"shape": [5, [1, 2], -1]}}' d = json.loads(s) @@ -20,22 +24,10 @@ def test_deserialize() -> None: async def test_encode() -> None: arr_flat = np.arange(60, dtype="uint8") arr_3x2 = np.reshape(arr_flat, (3, 4, 5)) - dt = UInt8() codec = ReshapeCodec(shape=((0, 1), 5)) proto = default_buffer_prototype() (encoded, *_) = await codec.encode( - [ - ( - proto.nd_buffer.from_numpy_array(arr_3x2), - ArraySpec( - arr_3x2.shape, - dt, - fill_value=0, - config=ArrayConfig("C", False), - prototype=proto, - ), - ) - ] + [(proto.nd_buffer.from_numpy_array(arr_3x2), array_spec(arr_3x2.shape))] ) assert encoded is not None assert encoded.shape == (12, 5) @@ -121,6 +113,17 @@ def test_validate_invalid(shape: tuple[int, ...], shape_spec: ShapeSpec) -> None ) +@pytest.mark.asyncio +async def test_decode_unexpected_shape() -> None: + codec = ReshapeCodec(shape=((0, 1),)) + buf = default_buffer_prototype().nd_buffer.from_numpy_array( + np.zeros(shape=(30, 40), dtype="uint8") + ) + spec = array_spec((3, 4)) + with pytest.raises(RuntimeError): + await codec.decode([(buf, spec)]) + + @pytest.mark.parametrize( ("input_shape", "shape_spec", "output_shape"), [ @@ -135,9 +138,7 @@ def test_valid( codec = ReshapeCodec(shape=shape_spec) dtype = UInt8() codec.validate(input_shape, dtype, RegularChunkGridMetadata(chunk_shape=input_shape)) - input_spec = ArraySpec( - input_shape, dtype, 0, ArrayConfig("C", False), default_buffer_prototype() - ) + input_spec = array_spec(input_shape) output_spec = codec.resolve_metadata(input_spec) assert output_spec.dtype == input_spec.dtype