Enforce strict mypy type checking and additional ruff rules

.pre-commit-config.yaml

@@ -23,6 +23,13 @@ repos:
     # Run the formatter.
   - id: ruff-format
 
+# Type checking with mypy
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: v1.15.0
+  hooks:
+  - id: mypy
+    additional_dependencies: ['types-PyYAML', 'types-requests']
+
 # Find common spelling mistakes in comments and docstrings
 - repo: https://github.com/codespell-project/codespell
   rev: v2.4.2

atlite/__init__.py

@@ -33,16 +33,16 @@
 __version__ = version("atlite")
 # e.g. "0.17.0" # TODO, in the network structure it should use the dev version
 match = re.match(r"(\d+\.\d+(\.\d+)?)", __version__)
-assert match, f"Could not determine release_version of pypsa: {__version__}"
+assert match, f"Could not determine release_version of atlite: {__version__}"
 release_version = match.group(0)
 
 __all__ = [
-    Cutout,
-    ExclusionContainer,
-    compute_indicatormatrix,
-    regrid,
-    cspinstallations,
-    solarpanels,
-    windturbines,
-    __version__,
+    "Cutout",
+    "ExclusionContainer",
+    "compute_indicatormatrix",
+    "regrid",
+    "cspinstallations",
+    "solarpanels",
+    "windturbines",
+    "__version__",
 ]

atlite/_types.py

@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: Contributors to atlite <https://github.com/pypsa/atlite>
+#
+# SPDX-License-Identifier: MIT
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Literal, TypeAlias, TypedDict
+
+import geopandas as gpd
+import numpy as np
+import scipy.sparse as sp
+import xarray as xr
+from pyproj import CRS
+from shapely.geometry.base import BaseGeometry
+
+NDArray: TypeAlias = np.ndarray[Any, np.dtype[np.floating[Any]]]
+NDArrayInt: TypeAlias = np.ndarray[Any, np.dtype[np.signedinteger[Any]]]
+NDArrayBool: TypeAlias = np.ndarray[Any, np.dtype[np.bool_]]
+DataArray: TypeAlias = xr.DataArray
+Dataset: TypeAlias = xr.Dataset
+PathLike: TypeAlias = str | Path
+NumericArray: TypeAlias = NDArray | DataArray
+Number: TypeAlias = int | float | np.number[Any]
+GeoDataFrame: TypeAlias = gpd.GeoDataFrame
+GeoSeries: TypeAlias = gpd.GeoSeries
+Geometry: TypeAlias = BaseGeometry
+CrsLike: TypeAlias = str | int | CRS | dict[str, Any] | None
+SparseMatrix: TypeAlias = sp.lil_matrix | sp.csr_matrix
+
+TrackingType: TypeAlias = (
+    Literal["horizontal", "tilted_horizontal", "vertical", "dual"] | None
+)
+ClearskyModel: TypeAlias = Literal["simple", "enhanced"]
+TrigonModel: TypeAlias = Literal["simple", "perez"]
+IrradiationType: TypeAlias = Literal["total", "direct", "diffuse", "ground"]
+HeatPumpSource: TypeAlias = Literal["air", "soil"]
+OrientationName: TypeAlias = Literal["latitude_optimal", "constant", "latitude"]
+DataFormat: TypeAlias = Literal["grib", "netcdf"]
+
+
+class ERA5RetrievalParams(TypedDict, total=False):
+    product: str
+    area: list[float]
+    grid: str
+    chunks: dict[str, int] | None
+    tmpdir: str | Path | None
+    lock: Any | None
+    data_format: Literal["grib", "netcdf"]
+    year: list[str]
+    month: list[str] | str
+    day: list[str] | str
+    time: str | list[str]
+    variable: str | list[str]

atlite/aggregate.py

@@ -1,29 +1,58 @@
 # SPDX-FileCopyrightText: Contributors to atlite <https://github.com/pypsa/atlite>
 #
 # SPDX-License-Identifier: MIT
-"""
-Functions for aggregating results.
-"""
+"""Functions for aggregating results."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
 
 import dask
 import xarray as xr
 
+if TYPE_CHECKING:
+    import pandas as pd
+    from scipy.sparse import spmatrix
+
+    from atlite._types import DataArray
+
+
+def aggregate_matrix(
+    da: DataArray,
+    matrix: spmatrix,
+    index: pd.Index,
+) -> DataArray:
+    """
+    Aggregate spatial data with a sparse matrix.
+
+    Parameters
+    ----------
+    da : xarray.DataArray
+        DataArray with spatial dimensions ``y`` and ``x``.
+    matrix : scipy.sparse.spmatrix
+        Aggregation matrix mapping flattened spatial cells to ``index``.
+    index : pandas.Index
+        Index defining the aggregated dimension.
 
-def aggregate_matrix(da, matrix, index):
+    Returns
+    -------
+    xarray.DataArray
+        Aggregated data indexed by ``index`` and, if present, time.
+    """
     if index.name is None:
         index = index.rename("dim_0")
     if isinstance(da.data, dask.array.core.Array):
         da = da.stack(spatial=("y", "x"))
-        da = da.chunk(dict(spatial=-1))
-        return xr.apply_ufunc(
+        da = da.chunk({"spatial": -1})
+        result = xr.apply_ufunc(
             lambda da: da * matrix.T,
             da,
             input_core_dims=[["spatial"]],
             output_core_dims=[[index.name]],
             dask="parallelized",
             output_dtypes=[da.dtype],
-            dask_gufunc_kwargs=dict(output_sizes={index.name: index.size}),
+            dask_gufunc_kwargs={"output_sizes": {index.name: index.size}},
         ).assign_coords(**{index.name: index})
-    else:
-        da = da.stack(spatial=("y", "x")).transpose("spatial", "time")
-        return xr.DataArray(matrix * da, [index, da.coords["time"]])
+        return cast("DataArray", result)
+    da = da.stack(spatial=("y", "x")).transpose("spatial", "time")
+    return xr.DataArray(matrix * da, [index, da.coords["time"]])