Add type hint checking with mypy (#460)

* Add type hint checking with mypy

* Fix EntryPoints import

* Fix cachedprotocolmeta typing

* Add release notes

* Run ruff-format

* pre-commit fixes

* More mypy fixes

* Fix release not location

* Remove type checking guards

* Some zarr3 typing fixes

* Fix inheritance

* Final zarr3 typing fixes

Author: dstansby

.pre-commit-config.yaml

@@ -25,3 +25,10 @@ repos:
     rev: 2024.08.19
     hooks:
     - id: sp-repo-review
+
+  -   repo: https://github.com/pre-commit/mirrors-mypy
+      rev: 'v1.4.0'
+      hooks:
+      -   id: mypy
+          args: [--config-file, pyproject.toml]
+          additional_dependencies: [numpy, pytest, zfpy, 'zarr==3.0.0b1']

docs/conf.py

@@ -12,7 +12,6 @@
 # All configuration values have a default; values that are commented out
 # serve to show the default.
 
-
 import os
 import sys
 from unittest.mock import Mock as MagicMock
@@ -232,7 +231,7 @@ def __getattr__(cls, name):
 
 # -- Options for LaTeX output ---------------------------------------------
 
-latex_elements = {
+latex_elements: dict[str, str] = {
     # The paper size ('letterpaper' or 'a4paper').
     #'papersize': 'letterpaper',
     # The font size ('10pt', '11pt' or '12pt').

docs/release.rst

@@ -26,6 +26,8 @@ Enhancements
   By :user:`Norman Rzepka <normanrz>`, :issue:`613`.
 * Add codec wrappers for Zarr 3.
   By :user:`Norman Rzepka <normanrz>`, :issue:`524`
+* Added mypy type checking to continuous integration.
+  By :user:`David Stansby <dstansby>`, :issue:`460`.
 
 Maintenance
 ~~~~~~~~~~~

numcodecs/abc.py

@@ -29,13 +29,14 @@
 """
 
 from abc import ABC, abstractmethod
+from typing import Optional
 
 
 class Codec(ABC):
     """Codec abstract base class."""
 
     # override in sub-class
-    codec_id = None
+    codec_id: Optional[str] = None
     """Codec identifier."""
 
     @abstractmethod

numcodecs/checksum32.py

@@ -1,19 +1,23 @@
 import struct
 import zlib
-from typing import Literal
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Literal
 
 import numpy as np
 
 from .abc import Codec
 from .compat import ensure_contiguous_ndarray, ndarray_copy
 from .jenkins import jenkins_lookup3
 
+if TYPE_CHECKING:
+    from typing_extensions import Buffer
+
 CHECKSUM_LOCATION = Literal['start', 'end']
 
 
 class Checksum32(Codec):
     # override in sub-class
-    checksum = None
+    checksum: Callable[["Buffer", int], int] | None = None
     location: CHECKSUM_LOCATION = 'start'
 
     def __init__(self, location: CHECKSUM_LOCATION | None = None):

numcodecs/compat.py

@@ -100,7 +100,7 @@ def ensure_contiguous_ndarray_like(buf, max_buffer_size=None, flatten=True) -> N
 
     # check for datetime or timedelta ndarray, the buffer interface doesn't support those
     if arr.dtype.kind in "Mm":
-        arr = arr.view(np.int64)
+        arr = arr.view(np.int64)  # type: ignore[arg-type]
 
     # check memory is contiguous, if so flatten
     if arr.flags.c_contiguous or arr.flags.f_contiguous:
@@ -117,7 +117,7 @@ def ensure_contiguous_ndarray_like(buf, max_buffer_size=None, flatten=True) -> N
     return arr
 
 
-def ensure_contiguous_ndarray(buf, max_buffer_size=None, flatten=True) -> np.array:
+def ensure_contiguous_ndarray(buf, max_buffer_size=None, flatten=True) -> np.ndarray:
     """Convenience function to coerce `buf` to a numpy array, if it is not already a
     numpy array. Also ensures that the returned value exports fully contiguous memory,
     and supports the new-style buffer interface. If the optional max_buffer_size is

numcodecs/lzma.py

@@ -1,11 +1,14 @@
-import contextlib
+from types import ModuleType
+from typing import Optional
 
-_lzma = None
+_lzma: Optional[ModuleType] = None
 try:
     import lzma as _lzma
 except ImportError:  # pragma: no cover
-    with contextlib.suppress(ImportError):
-        from backports import lzma as _lzma
+    try:
+        from backports import lzma as _lzma  # type: ignore[no-redef]
+    except ImportError:
+        pass
 
 
 if _lzma:

numcodecs/ndarray_like.py

@@ -1,7 +1,7 @@
 from typing import Any, ClassVar, Optional, Protocol, runtime_checkable
 
 
-class _CachedProtocolMeta(Protocol.__class__):
+class _CachedProtocolMeta(Protocol.__class__):  # type: ignore[name-defined]
     """Custom implementation of @runtime_checkable
 
     The native implementation of @runtime_checkable is slow,

numcodecs/registry.py

@@ -2,11 +2,13 @@
 applications to dynamically register and look-up codec classes."""
 
 import logging
-from importlib.metadata import entry_points
+from importlib.metadata import EntryPoints, entry_points
+
+from numcodecs.abc import Codec
 
 logger = logging.getLogger("numcodecs")
-codec_registry = {}
-entries = {}
+codec_registry: dict[str, Codec] = {}
+entries: dict[str, "EntryPoints"] = {}
 
 
 def run_entrypoints():

numcodecs/tests/test_lzma.py

@@ -1,5 +1,7 @@
 import itertools
 import unittest
+from types import ModuleType
+from typing import cast
 
 import numpy as np
 import pytest
@@ -20,6 +22,8 @@
     check_repr,
 )
 
+_lzma = cast(ModuleType, _lzma)
+
 codecs = [
     LZMA(),
     LZMA(preset=1),

numcodecs/tests/test_zarr3.py

@@ -1,11 +1,19 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import numpy as np
 import pytest
 
-zarr = pytest.importorskip("zarr")
+if not TYPE_CHECKING:
+    zarr = pytest.importorskip("zarr")
+else:
+    import zarr
+
+import zarr.storage
+from zarr.core.common import JSON
 
-import numcodecs.zarr3  # noqa: E402
+import numcodecs.zarr3
 
 pytestmark = [
     pytest.mark.skipif(zarr.__version__ < "3.0.0", reason="zarr 3.0.0 or later is required"),
@@ -17,7 +25,6 @@
 
 get_codec_class = zarr.registry.get_codec_class
 Array = zarr.Array
-JSON = zarr.core.common.JSON
 BytesCodec = zarr.codecs.BytesCodec
 Store = zarr.abc.store.Store
 MemoryStore = zarr.storage.MemoryStore
@@ -28,7 +35,7 @@
 
 
 @pytest.fixture
-def store() -> Store:
+def store() -> StorePath:
     return StorePath(MemoryStore(mode="w"))
 
 
@@ -43,6 +50,8 @@ def test_entry_points(codec_class: type[numcodecs.zarr3._NumcodecsCodec]):
 
 @pytest.mark.parametrize("codec_class", ALL_CODECS)
 def test_docstring(codec_class: type[numcodecs.zarr3._NumcodecsCodec]):
+    if codec_class.__doc__ is None:
+        pytest.skip()
     assert "See :class:`numcodecs." in codec_class.__doc__
 
 
@@ -59,7 +68,7 @@ def test_docstring(codec_class: type[numcodecs.zarr3._NumcodecsCodec]):
         numcodecs.zarr3.Shuffle,
     ],
 )
-def test_generic_codec_class(store: Store, codec_class: type[numcodecs.zarr3._NumcodecsCodec]):
+def test_generic_codec_class(store: StorePath, codec_class: type[numcodecs.zarr3._NumcodecsCodec]):
     data = np.arange(0, 256, dtype="uint16").reshape((16, 16))
 
     with pytest.warns(UserWarning, match=EXPECTED_WARNING_STR):
@@ -92,7 +101,9 @@ def test_generic_codec_class(store: Store, codec_class: type[numcodecs.zarr3._Nu
     ],
 )
 def test_generic_filter(
-    store: Store, codec_class: type[numcodecs.zarr3._NumcodecsCodec], codec_config: dict[str, JSON]
+    store: StorePath,
+    codec_class: type[numcodecs.zarr3._NumcodecsCodec],
+    codec_config: dict[str, JSON],
 ):
     data = np.linspace(0, 10, 256, dtype="float32").reshape((16, 16))
 
@@ -114,7 +125,7 @@ def test_generic_filter(
     np.testing.assert_array_equal(data, a[:, :])
 
 
-def test_generic_filter_bitround(store: Store):
+def test_generic_filter_bitround(store: StorePath):
     data = np.linspace(0, 1, 256, dtype="float32").reshape((16, 16))
 
     with pytest.warns(UserWarning, match=EXPECTED_WARNING_STR):
@@ -132,7 +143,7 @@ def test_generic_filter_bitround(store: Store):
     assert np.allclose(data, a[:, :], atol=0.1)
 
 
-def test_generic_filter_quantize(store: Store):
+def test_generic_filter_quantize(store: StorePath):
     data = np.linspace(0, 10, 256, dtype="float32").reshape((16, 16))
 
     with pytest.warns(UserWarning, match=EXPECTED_WARNING_STR):
@@ -150,7 +161,7 @@ def test_generic_filter_quantize(store: Store):
     assert np.allclose(data, a[:, :], atol=0.001)
 
 
-def test_generic_filter_packbits(store: Store):
+def test_generic_filter_packbits(store: StorePath):
     data = np.zeros((16, 16), dtype="bool")
     data[0:4, :] = True
 
@@ -189,7 +200,7 @@ def test_generic_filter_packbits(store: Store):
         numcodecs.zarr3.JenkinsLookup3,
     ],
 )
-def test_generic_checksum(store: Store, codec_class: type[numcodecs.zarr3._NumcodecsCodec]):
+def test_generic_checksum(store: StorePath, codec_class: type[numcodecs.zarr3._NumcodecsCodec]):
     data = np.linspace(0, 10, 256, dtype="float32").reshape((16, 16))
 
     with pytest.warns(UserWarning, match=EXPECTED_WARNING_STR):
@@ -208,7 +219,7 @@ def test_generic_checksum(store: Store, codec_class: type[numcodecs.zarr3._Numco
 
 
 @pytest.mark.parametrize("codec_class", [numcodecs.zarr3.PCodec, numcodecs.zarr3.ZFPY])
-def test_generic_bytes_codec(store: Store, codec_class: type[numcodecs.zarr3._NumcodecsCodec]):
+def test_generic_bytes_codec(store: StorePath, codec_class: type[numcodecs.zarr3._NumcodecsCodec]):
     try:
         codec_class()._codec  # noqa: B018
     except ValueError as e:

numcodecs/tests/test_zfpy.py

@@ -1,3 +1,6 @@
+from types import ModuleType
+from typing import cast
+
 import numpy as np
 import pytest
 
@@ -17,6 +20,9 @@
     check_repr,
 )
 
+_zfpy = cast(ModuleType, _zfpy)
+
+
 codecs = [
     ZFPY(mode=_zfpy.mode_fixed_rate, rate=-1),
     ZFPY(),