diff --git a/imod/common/interfaces/imodel.py b/imod/common/interfaces/imodel.py index 6398a9869..2fe91e0a4 100644 --- a/imod/common/interfaces/imodel.py +++ b/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 diff --git a/imod/common/serializer/netcdfserializer.py b/imod/common/serializer/netcdfserializer.py index e157818e6..de2e792f6 100644 --- a/imod/common/serializer/netcdfserializer.py +++ b/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 diff --git a/imod/common/serializer/zarrserializer.py b/imod/common/serializer/zarrserializer.py index 58ca34b48..86ad3753e 100644 --- a/imod/common/serializer/zarrserializer.py +++ b/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") diff --git a/imod/common/utilities/clip.py b/imod/common/utilities/clip.py index da1862a5c..7f78db73f 100644 --- a/imod/common/utilities/clip.py +++ b/imod/common/utilities/clip.py @@ -97,7 +97,7 @@ 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) @@ -105,7 +105,7 @@ def clip_by_grid(package: ILineDataPackage, active: xr.DataArray) -> ILineDataPa # 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: diff --git a/imod/common/utilities/dataclass_type.py b/imod/common/utilities/dataclass_type.py index 1fa53c515..86f6bf132 100644 --- a/imod/common/utilities/dataclass_type.py +++ b/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) diff --git a/imod/common/utilities/dtype.py b/imod/common/utilities/dtype.py index c2482919d..baf20ecf4 100644 --- a/imod/common/utilities/dtype.py +++ b/imod/common/utilities/dtype.py @@ -5,11 +5,17 @@ """ 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) @@ -17,8 +23,9 @@ def is_float(dtype: DTypeLike) -> bool: # 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 ) @@ -26,11 +33,13 @@ 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) diff --git a/imod/common/utilities/dump_model.py b/imod/common/utilities/dump_model.py index bdcf25de9..b94d5a02f 100644 --- a/imod/common/utilities/dump_model.py +++ b/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( diff --git a/imod/common/utilities/mask.py b/imod/common/utilities/mask.py index acfdeb77d..21f689f27 100644 --- a/imod/common/utilities/mask.py +++ b/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) diff --git a/imod/common/utilities/regrid.py b/imod/common/utilities/regrid.py index a3d5995b3..ca1848abe 100644 --- a/imod/common/utilities/regrid.py +++ b/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) } diff --git a/imod/common/utilities/schemata.py b/imod/common/utilities/schemata.py index 2060362f7..8a61aac87 100644 --- a/imod/common/utilities/schemata.py +++ b/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': []}`` """ - 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 diff --git a/imod/common/utilities/value_filters.py b/imod/common/utilities/value_filters.py index f3d4cbd4c..a3a6b03e8 100644 --- a/imod/common/utilities/value_filters.py +++ b/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: diff --git a/imod/msw/model.py b/imod/msw/model.py index e25ec39fd..f295b76d2 100644 --- a/imod/msw/model.py +++ b/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", diff --git a/pyproject.toml b/pyproject.toml index 2910acb20..366b87efe 100644 --- a/pyproject.toml +++ b/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.*",