normalize in codec_pipeline

Author: dcherian

src/zarr/core/codec_pipeline.py

@@ -35,6 +35,21 @@
 U = TypeVar("U")
 
 
+def normalize_slices(
+    idxr: tuple[int | slice, ...], shape: tuple[int, ...]
+) -> tuple[int | slice, ...]:
+    # replace slice objects with stop==None with size
+    out = []
+    for i, size in zip(idxr, shape, strict=False):
+        if not isinstance(i, slice):
+            out.append(i)
+            continue
+        if i.step not in [1, None]:
+            out.append(i)
+        out.append(slice(i.start, i.stop if i.stop is not None else size, i.step))
+    return tuple(out)
+
+
 def _unzip2(iterable: Iterable[tuple[T, U]]) -> tuple[list[T], list[U]]:
     out0: list[T] = []
     out1: list[U] = []
@@ -322,7 +337,9 @@ def _merge_chunk_array(
                     for idx in range(chunk_spec.ndim)
                 )
                 chunk_value = chunk_value[item]
-        chunk_array[chunk_selection] = chunk_value
+
+        normalized_selection = normalize_slices(chunk_selection, chunk_array.shape)
+        chunk_array[normalized_selection] = chunk_value
         return chunk_array
 
     async def write_batch(