lint

Author: d-v-b

src/zarr/codecs/sharding.py

@@ -544,7 +544,9 @@ def _encode_shard_dict_sync(
         if self.index_location == ShardingCodecIndexLocation.start:
             empty_chunks_mask = index.offsets_and_lengths[..., 0] == MAX_UINT_64
             index.offsets_and_lengths[~empty_chunks_mask, 0] += len(index_bytes)
-            index_bytes = self._encode_shard_index_sync(index)  # encode again with corrected offsets
+            index_bytes = self._encode_shard_index_sync(
+                index
+            )  # encode again with corrected offsets
             buffers.insert(0, index_bytes)
         else:
             buffers.append(index_bytes)
@@ -966,10 +968,10 @@ def _decode_shard_index_sync(
             spec = bb.resolve_metadata(spec)
 
         # Decode: reverse bb, then ab, then reverse aa
-        chunk_bytes: Buffer | None = index_bytes
+        chunk_bytes: Buffer = index_bytes
         for bb_codec, s in reversed(bb_with_spec):
             chunk_bytes = bb_codec._decode_sync(chunk_bytes, s)
-        chunk_array = ab_codec._decode_sync(chunk_bytes, ab_spec)
+        chunk_array: NDBuffer = ab_codec._decode_sync(chunk_bytes, ab_spec)
         for aa_codec, s in reversed(aa_with_spec):
             chunk_array = aa_codec._decode_sync(chunk_array, s)
 
@@ -984,18 +986,19 @@ def _encode_shard_index_sync(self, index: _ShardIndex) -> Buffer:
 
         aa_codecs, ab_codec, bb_codecs = codecs_from_list(list(self.index_codecs))
 
-        chunk_array: NDBuffer | None = get_ndbuffer_class().from_numpy_array(
-            index.offsets_and_lengths
-        )
+        aa_out: NDBuffer | None = get_ndbuffer_class().from_numpy_array(index.offsets_and_lengths)
 
         # Encode: aa forward, then ab, then bb forward
         spec = index_chunk_spec
         for aa_codec in aa_codecs:
-            chunk_array = aa_codec._encode_sync(chunk_array, spec)
+            assert aa_out is not None
+            aa_out = aa_codec._encode_sync(aa_out, spec)
             spec = aa_codec.resolve_metadata(spec)
-        chunk_bytes = ab_codec._encode_sync(chunk_array, spec)
+        assert aa_out is not None
+        chunk_bytes = ab_codec._encode_sync(aa_out, spec)
         spec = ab_codec.resolve_metadata(spec)
         for bb_codec in bb_codecs:
+            assert chunk_bytes is not None
             chunk_bytes = bb_codec._encode_sync(chunk_bytes, spec)
             spec = bb_codec.resolve_metadata(spec)
 

src/zarr/core/array.py

@@ -1993,9 +1993,8 @@ def _can_use_sync_path(self) -> bool:
         """
         pipeline = self.async_array.codec_pipeline
         store_path = self.async_array.store_path
-        return (
-            getattr(pipeline, "supports_sync_io", False)
-            and getattr(store_path, "supports_sync", False)
+        return getattr(pipeline, "supports_sync_io", False) and getattr(
+            store_path, "supports_sync", False
         )
 
     @classmethod

src/zarr/core/codec_pipeline.py

@@ -322,15 +322,22 @@ def _encode_one(
             return None
 
         spec = chunk_spec
+        aa_out: NDBuffer | None = chunk_array
 
         for aa_codec in self.array_array_codecs:
-            chunk_array = aa_codec._encode_sync(chunk_array, spec)
+            if aa_out is None:
+                return None
+            aa_out = aa_codec._encode_sync(aa_out, spec)
             spec = aa_codec.resolve_metadata(spec)
 
-        chunk_bytes = self.array_bytes_codec._encode_sync(chunk_array, spec)
+        if aa_out is None:
+            return None
+        chunk_bytes = self.array_bytes_codec._encode_sync(aa_out, spec)
         spec = self.array_bytes_codec.resolve_metadata(spec)
 
         for bb_codec in self.bytes_bytes_codecs:
+            if chunk_bytes is None:
+                return None
             chunk_bytes = bb_codec._encode_sync(chunk_bytes, spec)
             spec = bb_codec.resolve_metadata(spec)
 
@@ -694,7 +701,7 @@ async def _write_batch_compute(
     def _merge_and_filter(
         self,
         chunk_array_decoded: Iterable[NDBuffer | None],
-        batch_info: list,
+        batch_info: list[tuple[Any, ArraySpec, SelectorTuple, SelectorTuple, bool]],
         value: NDBuffer,
         drop_axes: tuple[int, ...],
     ) -> list[NDBuffer | None]:
@@ -717,9 +724,7 @@ def _merge_and_filter(
             ) in zip(chunk_array_decoded, batch_info, strict=False)
         ]
         chunk_array_batch: list[NDBuffer | None] = []
-        for chunk_array, (_, chunk_spec, *_) in zip(
-            chunk_array_merged, batch_info, strict=False
-        ):
+        for chunk_array, (_, chunk_spec, *_) in zip(chunk_array_merged, batch_info, strict=False):
             if chunk_array is None:
                 chunk_array_batch.append(None)  # type: ignore[unreachable]
             else:

tests/test_sync_codec_pipeline.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+from typing import Any
+
 import numpy as np
 import pytest
 
@@ -17,9 +19,7 @@
 from zarr.storage import MemoryStore
 
 
-def _make_array_spec(
-    shape: tuple[int, ...], dtype: np.dtype
-) -> ArraySpec:
+def _make_array_spec(shape: tuple[int, ...], dtype: np.dtype[Any]) -> ArraySpec:
     zdtype = get_data_type_from_native_dtype(dtype)
     return ArraySpec(
         shape=shape,
@@ -30,7 +30,7 @@ def _make_array_spec(
     )
 
 
-def _make_nd_buffer(arr: np.ndarray) -> zarr.core.buffer.NDBuffer:
+def _make_nd_buffer(arr: np.ndarray[Any, Any]) -> zarr.core.buffer.NDBuffer:
     return default_buffer_prototype().nd_buffer.from_numpy_array(arr)
 
 
@@ -40,19 +40,19 @@ def _make_nd_buffer(arr: np.ndarray) -> zarr.core.buffer.NDBuffer:
 
 
 class TestSupportsSync:
-    def test_gzip_supports_sync(self):
+    def test_gzip_supports_sync(self) -> None:
         assert GzipCodec().supports_sync
 
-    def test_zstd_supports_sync(self):
+    def test_zstd_supports_sync(self) -> None:
         assert ZstdCodec().supports_sync
 
-    def test_bytes_supports_sync(self):
+    def test_bytes_supports_sync(self) -> None:
         assert BytesCodec().supports_sync
 
-    def test_transpose_supports_sync(self):
+    def test_transpose_supports_sync(self) -> None:
         assert TransposeCodec(order=(0, 1)).supports_sync
 
-    def test_sharding_supports_sync(self):
+    def test_sharding_supports_sync(self) -> None:
         from zarr.codecs.sharding import ShardingCodec
 
         assert ShardingCodec(chunk_shape=(8,)).supports_sync
@@ -64,7 +64,7 @@ def test_sharding_supports_sync(self):
 
 
 class TestGzipCodecSync:
-    def test_roundtrip(self):
+    def test_roundtrip(self) -> None:
         codec = GzipCodec(level=1)
         arr = np.arange(100, dtype="float64")
         spec = _make_array_spec(arr.shape, arr.dtype)
@@ -78,7 +78,7 @@ def test_roundtrip(self):
 
 
 class TestZstdCodecSync:
-    def test_roundtrip(self):
+    def test_roundtrip(self) -> None:
         codec = ZstdCodec(level=1)
         arr = np.arange(100, dtype="float64")
         spec = _make_array_spec(arr.shape, arr.dtype)
@@ -92,7 +92,7 @@ def test_roundtrip(self):
 
 
 class TestBytesCodecSync:
-    def test_roundtrip(self):
+    def test_roundtrip(self) -> None:
         codec = BytesCodec()
         arr = np.arange(100, dtype="float64")
         spec = _make_array_spec(arr.shape, arr.dtype)
@@ -108,7 +108,7 @@ def test_roundtrip(self):
 
 
 class TestTransposeCodecSync:
-    def test_roundtrip(self):
+    def test_roundtrip(self) -> None:
         codec = TransposeCodec(order=(1, 0))
         arr = np.arange(12, dtype="float64").reshape(3, 4)
         spec = _make_array_spec(arr.shape, arr.dtype)
@@ -127,29 +127,31 @@ def test_roundtrip(self):
 
 
 class TestPipelineConstruction:
-    def test_from_codecs_valid(self):
+    def test_from_codecs_valid(self) -> None:
         pipeline = BatchedCodecPipeline.from_codecs([BytesCodec(), GzipCodec(level=1)])
         assert isinstance(pipeline, BatchedCodecPipeline)
         assert len(pipeline.bytes_bytes_codecs) == 1
         assert isinstance(pipeline.array_bytes_codec, BytesCodec)
 
-    def test_from_codecs_accepts_sharding(self):
+    def test_from_codecs_accepts_sharding(self) -> None:
         from zarr.codecs.sharding import ShardingCodec
 
         pipeline = BatchedCodecPipeline.from_codecs([ShardingCodec(chunk_shape=(8,))])
         assert isinstance(pipeline, BatchedCodecPipeline)
         assert pipeline._all_sync
 
-    def test_from_codecs_rejects_missing_array_bytes(self):
+    def test_from_codecs_rejects_missing_array_bytes(self) -> None:
         with pytest.raises(ValueError, match="Required ArrayBytesCodec"):
             BatchedCodecPipeline.from_codecs([GzipCodec()])
 
-    def test_from_codecs_with_transpose(self):
-        pipeline = BatchedCodecPipeline.from_codecs([
-            TransposeCodec(order=(1, 0)),
-            BytesCodec(),
-            GzipCodec(level=1),
-        ])
+    def test_from_codecs_with_transpose(self) -> None:
+        pipeline = BatchedCodecPipeline.from_codecs(
+            [
+                TransposeCodec(order=(1, 0)),
+                BytesCodec(),
+                GzipCodec(level=1),
+            ]
+        )
         assert len(pipeline.array_array_codecs) == 1
         assert isinstance(pipeline.array_array_codecs[0], TransposeCodec)
 
@@ -161,21 +163,21 @@ def test_from_codecs_with_transpose(self):
 
 class TestPipelineRoundtrip:
     @pytest.mark.asyncio
-    async def test_encode_decode_single_chunk(self):
+    async def test_encode_decode_single_chunk(self) -> None:
         pipeline = BatchedCodecPipeline.from_codecs([BytesCodec(), GzipCodec(level=1)])
         arr = np.random.default_rng(42).standard_normal((32, 32)).astype("float64")
         spec = _make_array_spec(arr.shape, arr.dtype)
         pipeline = pipeline.evolve_from_array_spec(spec)
         nd_buf = _make_nd_buffer(arr)
 
         encoded = await pipeline.encode([(nd_buf, spec)])
-        decoded = await pipeline.decode([(list(encoded)[0], spec)])
-        result = list(decoded)[0]
+        decoded = await pipeline.decode([(next(iter(encoded)), spec)])
+        result = next(iter(decoded))
         assert result is not None
         np.testing.assert_array_equal(arr, result.as_numpy_array())
 
     @pytest.mark.asyncio
-    async def test_encode_decode_multiple_chunks(self):
+    async def test_encode_decode_multiple_chunks(self) -> None:
         pipeline = BatchedCodecPipeline.from_codecs([BytesCodec(), GzipCodec(level=1)])
         rng = np.random.default_rng(42)
         spec = _make_array_spec((16, 16), np.dtype("float64"))
@@ -185,20 +187,20 @@ async def test_encode_decode_multiple_chunks(self):
 
         encoded = list(await pipeline.encode([(buf, spec) for buf in nd_bufs]))
         decoded = list(await pipeline.decode([(enc, spec) for enc in encoded]))
-        for original, dec in zip(chunks, decoded):
+        for original, dec in zip(chunks, decoded, strict=False):
             assert dec is not None
             np.testing.assert_array_equal(original, dec.as_numpy_array())
 
     @pytest.mark.asyncio
-    async def test_encode_decode_empty_batch(self):
+    async def test_encode_decode_empty_batch(self) -> None:
         pipeline = BatchedCodecPipeline.from_codecs([BytesCodec(), GzipCodec(level=1)])
         encoded = await pipeline.encode([])
         assert list(encoded) == []
         decoded = await pipeline.decode([])
         assert list(decoded) == []
 
     @pytest.mark.asyncio
-    async def test_encode_decode_none_chunk(self):
+    async def test_encode_decode_none_chunk(self) -> None:
         pipeline = BatchedCodecPipeline.from_codecs([BytesCodec(), GzipCodec(level=1)])
         spec = _make_array_spec((8,), np.dtype("float64"))
         pipeline = pipeline.evolve_from_array_spec(spec)
@@ -216,7 +218,7 @@ async def test_encode_decode_none_chunk(self):
 
 
 class TestDefaultPipelineSync:
-    def test_create_array_uses_batched_pipeline(self):
+    def test_create_array_uses_batched_pipeline(self) -> None:
         store = MemoryStore()
         arr = zarr.create_array(
             store,
@@ -230,7 +232,7 @@ def test_create_array_uses_batched_pipeline(self):
         arr[:] = data
         np.testing.assert_array_equal(arr[:], data)
 
-    def test_open_uses_batched_pipeline(self):
+    def test_open_uses_batched_pipeline(self) -> None:
         store = MemoryStore()
         arr = zarr.create_array(
             store,
@@ -245,7 +247,7 @@ def test_open_uses_batched_pipeline(self):
         assert isinstance(arr2.async_array.codec_pipeline, BatchedCodecPipeline)
         np.testing.assert_array_equal(arr2[:], data)
 
-    def test_from_array_uses_batched_pipeline(self):
+    def test_from_array_uses_batched_pipeline(self) -> None:
         store1 = MemoryStore()
         arr1 = zarr.create_array(
             store1,
@@ -261,7 +263,7 @@ def test_from_array_uses_batched_pipeline(self):
         assert isinstance(arr2.async_array.codec_pipeline, BatchedCodecPipeline)
         np.testing.assert_array_equal(arr2[:], data)
 
-    def test_partial_write(self):
+    def test_partial_write(self) -> None:
         store = MemoryStore()
         arr = zarr.create_array(
             store,
@@ -276,7 +278,7 @@ def test_partial_write(self):
         expected[5:15] = np.arange(10, dtype="int32") + 1
         np.testing.assert_array_equal(result, expected)
 
-    def test_zstd_codec(self):
+    def test_zstd_codec(self) -> None:
         store = MemoryStore()
         arr = zarr.create_array(
             store,
@@ -289,18 +291,16 @@ def test_zstd_codec(self):
         arr[:] = data
         np.testing.assert_array_equal(arr[:], data)
 
-    def test_supports_sync_io(self):
+    def test_supports_sync_io(self) -> None:
         """Default pipeline supports sync IO when all codecs are sync."""
         pipeline = BatchedCodecPipeline.from_codecs([BytesCodec(), GzipCodec(level=1)])
         assert pipeline.supports_sync_io
 
-    def test_config_switch_to_sync_pipeline_compat(self):
+    def test_config_switch_to_sync_pipeline_compat(self) -> None:
         """Verify backwards compat: SyncCodecPipeline config path still works."""
         from zarr.experimental.sync_codecs import SyncCodecPipeline
 
-        zarr.config.set(
-            {"codec_pipeline.path": "zarr.experimental.sync_codecs.SyncCodecPipeline"}
-        )
+        zarr.config.set({"codec_pipeline.path": "zarr.experimental.sync_codecs.SyncCodecPipeline"})
         try:
             store = MemoryStore()
             arr = zarr.create_array(store, shape=(10,), dtype="float64")