Fix some mypy errors (#3044)

* Clean mypy config

* Fix typing errors in entrypoint test package

* Fix test_transpose typing errors

* Fix typing errors in test_config

---------

Co-authored-by: Davis Bennett <[email protected]>

Author: dstansby

pyproject.toml

@@ -348,27 +348,27 @@ python_version = "3.11"
 ignore_missing_imports = true
 namespace_packages = false
 
-
 strict = true
 warn_unreachable = true
-
 enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"]
 
+
 [[tool.mypy.overrides]]
 module = [
-    "zarr.v2.*",
+    "tests.package_with_entrypoint.*",
+    "tests.test_codecs.test_transpose",
+    "tests.test_config"
 ]
-ignore_errors = true
+strict = false
 
+# TODO: Move the next modules up to the strict = false section
+# and fix the errors
 [[tool.mypy.overrides]]
 module = [
     "zarr.testing.stateful",  # lots of hypothesis decorator errors
-    "tests.package_with_entrypoint.*",
     "tests.test_codecs.test_codecs",
-    "tests.test_codecs.test_transpose",
     "tests.test_metadata.*",
     "tests.test_store.*",
-    "tests.test_config",
     "tests.test_group",
     "tests.test_indexing",
     "tests.test_properties",

tests/package_with_entrypoint/__init__.py

@@ -1,13 +1,14 @@
 from collections.abc import Iterable
+from typing import Any
 
-from numpy import ndarray
+import numpy as np
+import numpy.typing as npt
 
 import zarr.core.buffer
-from zarr.abc.codec import ArrayBytesCodec, CodecInput, CodecOutput, CodecPipeline
+from zarr.abc.codec import ArrayBytesCodec, CodecInput, CodecPipeline
 from zarr.codecs import BytesCodec
 from zarr.core.array_spec import ArraySpec
 from zarr.core.buffer import Buffer, NDBuffer
-from zarr.core.common import BytesLike
 
 
 class TestEntrypointCodec(ArrayBytesCodec):
@@ -16,14 +17,14 @@ class TestEntrypointCodec(ArrayBytesCodec):
     async def encode(
         self,
         chunks_and_specs: Iterable[tuple[CodecInput | None, ArraySpec]],
-    ) -> Iterable[CodecOutput | None]:
-        pass
+    ) -> Iterable[Buffer | None]:
+        return [None]
 
     async def decode(
         self,
         chunks_and_specs: Iterable[tuple[CodecInput | None, ArraySpec]],
-    ) -> ndarray:
-        pass
+    ) -> npt.NDArray[Any]:
+        return np.array(1)
 
     def compute_encoded_size(self, input_byte_length: int, chunk_spec: ArraySpec) -> int:
         return input_byte_length
@@ -35,13 +36,13 @@ def __init__(self, batch_size: int = 1) -> None:
 
     async def encode(
         self, chunks_and_specs: Iterable[tuple[CodecInput | None, ArraySpec]]
-    ) -> BytesLike:
-        pass
+    ) -> Iterable[Buffer | None]:
+        return [None]
 
     async def decode(
         self, chunks_and_specs: Iterable[tuple[CodecInput | None, ArraySpec]]
-    ) -> ndarray:
-        pass
+    ) -> Iterable[NDBuffer | None]:
+        return np.array(1)
 
 
 class TestEntrypointBuffer(Buffer):

tests/test_codecs/test_transpose.py

@@ -48,6 +48,7 @@ async def test_transpose(
     read_data = await _AsyncArrayProxy(a)[:, :].get()
     assert np.array_equal(data, read_data)
 
+    assert isinstance(read_data, np.ndarray)
     if runtime_read_order == "F":
         assert read_data.flags["F_CONTIGUOUS"]
         assert not read_data.flags["C_CONTIGUOUS"]
@@ -90,5 +91,5 @@ def test_transpose_invalid(
                 dtype=data.dtype,
                 fill_value=0,
                 chunk_key_encoding={"name": "v2", "separator": "."},
-                filters=[TransposeCodec(order=order)],
+                filters=[TransposeCodec(order=order)],  # type: ignore[arg-type]
             )

tests/test_config.py

@@ -10,7 +10,7 @@
 import zarr
 import zarr.api
 from zarr import zeros
-from zarr.abc.codec import CodecInput, CodecOutput, CodecPipeline
+from zarr.abc.codec import CodecPipeline
 from zarr.abc.store import ByteSetter, Store
 from zarr.codecs import (
     BloscCodec,
@@ -21,6 +21,7 @@
 )
 from zarr.core.array_spec import ArraySpec
 from zarr.core.buffer import NDBuffer
+from zarr.core.buffer.core import Buffer
 from zarr.core.codec_pipeline import BatchedCodecPipeline
 from zarr.core.config import BadConfigError, config
 from zarr.core.indexing import SelectorTuple
@@ -144,7 +145,7 @@ def test_config_codec_pipeline_class(store: Store) -> None:
     class MockCodecPipeline(BatchedCodecPipeline):
         async def write(
             self,
-            batch_info: Iterable[tuple[ByteSetter, ArraySpec, SelectorTuple, SelectorTuple]],
+            batch_info: Iterable[tuple[ByteSetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],
             value: NDBuffer,
             drop_axes: tuple[int, ...] = (),
         ) -> None:
@@ -174,7 +175,7 @@ async def write(
     class MockEnvCodecPipeline(CodecPipeline):
         pass
 
-    register_pipeline(MockEnvCodecPipeline)
+    register_pipeline(MockEnvCodecPipeline)  # type: ignore[type-abstract]
 
     with mock.patch.dict(
         os.environ, {"ZARR_CODEC_PIPELINE__PATH": fully_qualified_name(MockEnvCodecPipeline)}
@@ -191,10 +192,9 @@ def test_config_codec_implementation(store: Store) -> None:
     _mock = Mock()
 
     class MockBloscCodec(BloscCodec):
-        async def _encode_single(
-            self, chunk_data: CodecInput, chunk_spec: ArraySpec
-        ) -> CodecOutput | None:
+        async def _encode_single(self, chunk_bytes: Buffer, chunk_spec: ArraySpec) -> Buffer | None:
             _mock.call()
+            return None
 
     register_codec("blosc", MockBloscCodec)
     with config.set({"codecs.blosc": fully_qualified_name(MockBloscCodec)}):
@@ -245,7 +245,7 @@ def test_config_buffer_implementation() -> None:
     # has default value
     assert fully_qualified_name(get_buffer_class()) == config.defaults[0]["buffer"]
 
-    arr = zeros(shape=(100), store=StoreExpectingTestBuffer())
+    arr = zeros(shape=(100,), store=StoreExpectingTestBuffer())
 
     # AssertionError of StoreExpectingTestBuffer when not using my buffer
     with pytest.raises(AssertionError):