Issue #1877 type hints common module

imod/common/interfaces/imodel.py

@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from typing import Optional, Tuple
+from typing import Any, Optional, Tuple
 
 from imod.common.interfaces.idict import IDict
 from imod.common.statusinfo import StatusInfoBase
@@ -37,7 +37,7 @@ def domain(self):
 
     @property
     @abstractmethod
-    def options(self) -> dict:
+    def options(self) -> dict[str, Any]:
         raise NotImplementedError
 
     @property

imod/common/serializer/netcdfserializer.py

@@ -59,7 +59,7 @@ def _dataset_encoding(self, pkg: IPackageBase) -> GridDataset:
 
         return pkg.dataset
 
-    def _netcdf_encoding(self, pkg: IPackageBase) -> dict:
+    def _netcdf_encoding(self, pkg: IPackageBase) -> dict[str, dict[str, str]]:
         """
 
         The encoding used in the to_netcdf method

imod/common/serializer/zarrserializer.py

@@ -2,7 +2,7 @@
 from contextlib import nullcontext
 from copy import deepcopy
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import xugrid as xu
 
@@ -28,6 +28,9 @@ def __init__(self, use_zip: bool = False):
     def to_file(
         self, pkg: IPackageBase, directory: Path, file_name: str, **kwargs
     ) -> Path:
+        store: zarr.storage.ZipStore | Path
+        write_context: zarr.storage.ZipStore | nullcontext[Any]
+
         if self.use_zip:
             path = directory / f"{file_name}.zarr.zip"
             store = zarr.storage.ZipStore(path, mode="w")

imod/common/utilities/clip.py

@@ -97,15 +97,15 @@ def clip_by_grid(  # noqa: F811
     return cls._from_dataset(selection)
 
 
-@dispatch  # type: ignore[no-redef, misc]
+@dispatch  # type: ignore[no-redef]
 def clip_by_grid(package: ILineDataPackage, active: xr.DataArray) -> ILineDataPackage:  # noqa: F811
     """Clip LineDataPackage outside structured grid."""
     return _clip_by_grid_line_data(package, active)
 
 
 # For some reason the plum dispatching finds "active: GridDataArray" ambiguous
 # and raises an error if this is not duplicated.
-@dispatch  # type: ignore[no-redef, misc]
+@dispatch  # type: ignore[no-redef]
 def clip_by_grid(  # noqa: F811
     package: ILineDataPackage, active: xu.UgridDataArray
 ) -> ILineDataPackage:

imod/common/utilities/dataclass_type.py

@@ -1,4 +1,4 @@
-from typing import ClassVar, Protocol, runtime_checkable
+from typing import Any, ClassVar, Protocol, runtime_checkable
 
 from pydantic import ConfigDict
 from pydantic.dataclasses import dataclass
@@ -14,9 +14,9 @@ class DataclassType(Protocol):
     # See also:
     # https://github.com/python/mypy/issues/6568#issuecomment-1324196557
 
-    __dataclass_fields__: ClassVar[dict]
+    __dataclass_fields__: ClassVar[dict[str, Any]]
 
-    def asdict(self) -> dict:
+    def asdict(self) -> dict[str, Any]:
         return vars(self)
 
 

imod/common/utilities/dtype.py

@@ -5,32 +5,41 @@
 """
 
 import numbers
+from typing import Any, Protocol, cast
 
 import numpy as np
 from numpy.typing import DTypeLike
 
 
+class SupportsDTypeType(Protocol):
+    @property
+    def type(self) -> type[Any]: ...
+
+
 def is_float(dtype: DTypeLike) -> bool:
     try:
         return np.issubdtype(dtype, np.floating)
     except TypeError:
         # Catch cases where dtype is not a numpy dtype and check if subclass is
         # not an integer. As numpy-style integers are also considered real
         # numbers.
-        return issubclass(dtype.type, numbers.Real) and not issubclass(
-            dtype.type, numbers.Integral
+        dtype_with_type = cast(SupportsDTypeType, dtype)
+        return issubclass(dtype_with_type.type, numbers.Real) and not issubclass(
+            dtype_with_type.type, numbers.Integral
         )
 
 
 def is_integer(dtype: DTypeLike) -> bool:
     try:
         return np.issubdtype(dtype, np.integer)
     except TypeError:
-        return issubclass(dtype.type, numbers.Integral)
+        dtype_with_type = cast(SupportsDTypeType, dtype)
+        return issubclass(dtype_with_type.type, numbers.Integral)
 
 
 def is_bool(dtype: DTypeLike) -> bool:
     try:
         return np.issubdtype(dtype, np.bool_)
     except TypeError:
-        return issubclass(dtype.type, np.bool)
+        dtype_with_type = cast(SupportsDTypeType, dtype)
+        return issubclass(dtype_with_type.type, np.bool)

imod/common/utilities/dump_model.py

@@ -1,6 +1,6 @@
 import collections
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, DefaultDict, Optional
 
 import tomli_w
 
@@ -16,7 +16,7 @@ def dump_model(
     model: IDict,
     directory,
     modelname,
-    validate: Optional[bool] = True,
+    validate: bool = True,
     mdal_compliant: bool = False,
     crs: Optional[Any] = None,
     engine: EngineType = "netcdf4",
@@ -64,7 +64,7 @@ def dump_model(
             if statusinfo.has_errors():
                 raise ValidationError(statusinfo.to_string())
 
-    toml_content: dict = collections.defaultdict(dict)
+    toml_content: DefaultDict[str, dict[str, Any]] = collections.defaultdict(dict)
 
     for pkgname, pkg in model.items():
         pkg_path = pkg.to_file(

imod/common/utilities/mask.py

@@ -77,6 +77,7 @@ def mask_package(package: IPackage, mask: GridDataArray) -> IPackage:
     masked = {}
 
     for var in package.dataset.data_vars.keys():
+        var = str(var)
         if _skip_dataarray(package.dataset[var]) or _skip_variable(package, var):
             masked[var] = package.dataset[var]
         else:
@@ -114,6 +115,7 @@ def mask_da(da: GridDataArray, mask: GridDataArray) -> GridDataArray:
     dtype of the original DataArray. It will set the
     value to 0 for integers, np.nan for floats, and False for booleans.
     """
+    other: int | float | bool
 
     if is_integer(da.dtype):
         other = MaskValues.integer
@@ -127,8 +129,11 @@ def mask_da(da: GridDataArray, mask: GridDataArray) -> GridDataArray:
         )
     # Align the mask, as calling where with "other" specified does not
     # automatically align the mask to the DataArray.
-    _, mask = xr.align(da, mask, join="left", copy=False)
-    return da.where(mask, other=other)
+    aligned: tuple[GridDataArray, GridDataArray] = xr.align(
+        da, mask, join="left", copy=False
+    )
+    _, mask_aligned = aligned
+    return da.where(mask_aligned, other=other)
 
 
 def _mask_spatial_var_pkg(
@@ -213,7 +218,7 @@ def broadcast_and_mask_arrays(
     # will result in no spatial grid.
     if not is_spatial_grid(broadcasted_arrays[0]):
         raise ValueError("One or more arrays need to be a spatial grid.")
-    broadcasted_arrays = dict(zip(arrays.keys(), broadcasted_arrays))
+    broadcasted_arrays_dict = dict(zip(arrays.keys(), broadcasted_arrays))
 
     # Mask arrays with np.nan values
-    return mask_arrays(broadcasted_arrays)
+    return mask_arrays(broadcasted_arrays_dict)

imod/common/utilities/regrid.py

@@ -1,7 +1,7 @@
 import copy
 from collections import defaultdict
 from dataclasses import asdict
-from typing import Any, Optional, Tuple, TypeAlias, Union
+from typing import Any, DefaultDict, Optional, Tuple, TypeAlias, Union
 
 import numpy as np
 import xarray as xr
@@ -163,14 +163,23 @@ def _regrid_package_data(
     return new_package_data
 
 
-def _get_unique_regridder_types(model: IModel) -> defaultdict[RegridderType, list[str]]:
+def __get_regrid_methods_as_dict(package: IRegridPackage) -> dict[str, RegridVarType]:
+    """
+    Returns the regrid methods of a package as a dictionary. Separated function
+    to set an ignore for mypy.
+    """
+    regrid_methods = package.get_regrid_methods()
+    return asdict(regrid_methods)  # type: ignore[arg-type]
+
+
+def _get_unique_regridder_types(model: IModel) -> DefaultDict[RegridderType, list[Any]]:
     """
     This function loops over the packages and  collects all regridder-types that are in use.
     """
-    methods: defaultdict = defaultdict(list)
+    methods: DefaultDict[RegridderType, list[Any]] = defaultdict(list)
     regrid_packages = [pkg for pkg in model.values() if isinstance(pkg, IRegridPackage)]
     regrid_packages_with_methods = {
-        pkg: asdict(pkg.get_regrid_methods()).items()  # type: ignore[union-attr]
+        pkg: __get_regrid_methods_as_dict(pkg).items()
         for pkg in regrid_packages
         if not isinstance(pkg.get_regrid_methods(), EmptyRegridMethod)
     }

imod/common/utilities/schemata.py

@@ -1,10 +1,10 @@
 from collections import defaultdict
-from collections.abc import Mapping
 from copy import deepcopy
 from typing import Any, Optional, Protocol
 
 from imod.common.statusinfo import NestedStatusInfo, StatusInfo, StatusInfoBase
 from imod.schemata import BaseSchema, SchemataDict, ValidationError
+from imod.typing import GridDataset
 
 
 def filter_schemata_dict(
@@ -33,7 +33,7 @@ def filter_schemata_dict(
     Prints ``{'stage': [<imod.schemata.AllNoDataSchema at 0x1b152b12aa0>]}``
     """
 
-    d = {}
+    d: SchemataDict = {}
     for key, schema_ls in schemata_dict.items():
         schema_match = [
             schema for schema in schema_ls if isinstance(schema, schema_types)
@@ -62,7 +62,7 @@ def concatenate_schemata_dicts(
 
 
 def validate_schemata_dict(
-    schemata: SchemataDict, data: Mapping, **kwargs: Any
+    schemata: SchemataDict, data: GridDataset | dict[str, Any], **kwargs: Any
 ) -> dict[str, list[ValidationError]]:
     """
     Validate a data mapping against a schemata dictionary. Returns a dictionary

imod/common/utilities/value_filters.py

@@ -46,7 +46,7 @@ def is_empty_dataarray(da: Any) -> bool:
 
 def get_scalar_variables(ds: GridDataset) -> list[str]:
     """Returns scalar variables in a dataset."""
-    return [var for var, arr in ds.variables.items() if is_scalar(arr)]
+    return [str(var) for var, arr in ds.variables.items() if is_scalar(arr)]
 
 
 def enforce_scalar(a: GridDataArray) -> Any:

imod/msw/model.py

@@ -426,7 +426,7 @@ def dump(
         self,
         directory: Union[str, Path],
         modelname: Optional[str] = None,
-        validate: Optional[bool] = True,
+        validate: bool = True,
         mdal_compliant: bool = False,
         crs: Optional[str] = None,
         engine: EngineType = "netcdf4",

pyproject.toml

@@ -116,7 +116,6 @@ warn_return_any = false
 # works with the vscode mypy extension
 [[tool.mypy.overrides]]
 module = [
-    "imod.common.*",
     "imod.data.*",
     "imod.evaluate.*",
     "imod.formats.*",