Enforce 100% type-hint coverage and align typing to PEP 585/604

.pre-commit-config.yaml

@@ -48,4 +48,5 @@ repos:
     hooks:
       - id: mypy
         additional_dependencies: [types-PyYAML, types-Pillow, types-tqdm]
+        files: ^neural_lam/
         description: Check for type errors

CHANGELOG.md

@@ -74,6 +74,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Maintenance
 
+- Enforce 100% type-hint coverage across `neural_lam/` via a `[tool.mypy]` gate (`disallow_untyped_defs`/`disallow_incomplete_defs`) and align all type annotations with PEP 585 and PEP 604 [\#673](https://github.com/mllam/neural-lam/pull/673) @GiGiKoneti
+
 - Add comprehensive type hints to GraphLAM in `neural_lam/models/step_predictors/graph/graph_lam.py` [\#669](https://github.com/mllam/neural-lam/pull/669) @GiGiKoneti
 
 - Add comprehensive type hints to ARForecaster in `neural_lam/models/forecasters/autoregressive.py` [\#663](https://github.com/mllam/neural-lam/pull/663) @GiGiKoneti

neural_lam/config.py

@@ -3,7 +3,7 @@
 # Standard library
 import dataclasses
 from pathlib import Path
-from typing import Dict, Union
+from typing import cast
 
 # Third-party
 import dataclass_wizard
@@ -35,7 +35,7 @@ class DatastoreSelection:
     kind: str
     config_path: str
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         """
         Validate that the selected datastore kind is implemented.
 
@@ -56,11 +56,11 @@ class ManualStateFeatureWeighting:
 
     Attributes
     ----------
-    weights : Dict[str, float]
+    weights : dict[str, float]
         Manual weights for the state features.
     """
 
-    weights: Dict[str, float]
+    weights: dict[str, float]
 
 
 @dataclasses.dataclass
@@ -80,14 +80,14 @@ class OutputClamping:
 
     Attributes
     ----------
-    lower : Dict[str, float]
+    lower : dict[str, float]
         The minimum value to clamp each output feature to.
-    upper : Dict[str, float]
+    upper : dict[str, float]
         The maximum value to clamp each output feature to.
     """
 
-    lower: Dict[str, float] = dataclasses.field(default_factory=dict)
-    upper: Dict[str, float] = dataclasses.field(default_factory=dict)
+    lower: dict[str, float] = dataclasses.field(default_factory=dict)
+    upper: dict[str, float] = dataclasses.field(default_factory=dict)
 
 
 @dataclasses.dataclass
@@ -97,8 +97,8 @@ class TrainingConfig:
 
     Attributes
     ----------
-    state_feature_weighting : Union[ManualStateFeatureWeighting,
-                                    UniformFeatureWeighting]
+    state_feature_weighting :
+        ManualStateFeatureWeighting | UniformFeatureWeighting
         The method to use for weighting the state features in the loss
         function. Defaults to uniform weighting (`UniformFeatureWeighting`, i.e.
         all features are weighted equally).
@@ -107,9 +107,9 @@ class TrainingConfig:
         Defaults to an empty ``OutputClamping`` (no clamping).
     """
 
-    state_feature_weighting: Union[
-        ManualStateFeatureWeighting, UniformFeatureWeighting
-    ] = dataclasses.field(default_factory=UniformFeatureWeighting)
+    state_feature_weighting: (
+        ManualStateFeatureWeighting | UniformFeatureWeighting
+    ) = dataclasses.field(default_factory=UniformFeatureWeighting)
 
     output_clamping: OutputClamping = dataclasses.field(
         default_factory=OutputClamping
@@ -174,7 +174,7 @@ class InvalidConfigError(Exception):
 
 def load_config_and_datastore(
     config_path: str,
-) -> tuple[NeuralLAMConfig, Union[MDPDatastore, NpyFilesDatastoreMEPS]]:
+) -> tuple[NeuralLAMConfig, MDPDatastore | NpyFilesDatastoreMEPS]:
     """
     Load the neural-lam configuration and the datastore specified in the
     configuration.
@@ -186,7 +186,7 @@ def load_config_and_datastore(
 
     Returns
     -------
-    tuple[NeuralLAMConfig, Union[MDPDatastore, NpyFilesDatastoreMEPS]]
+    tuple[NeuralLAMConfig, MDPDatastore | NpyFilesDatastoreMEPS]
         The Neural-LAM configuration and the loaded datastore.
     """
     try:
@@ -204,4 +204,4 @@ def load_config_and_datastore(
         datastore_kind=config.datastore.kind, config_path=datastore_config_path
     )
 
-    return config, datastore
+    return config, cast(MDPDatastore | NpyFilesDatastoreMEPS, datastore)

neural_lam/create_graph.py

@@ -3,7 +3,6 @@
 # Standard library
 import os
 from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
-from typing import Optional
 
 # Third-party
 import matplotlib
@@ -22,7 +21,7 @@
 
 
 def plot_graph(
-    graph: pyg.data.Data, title: Optional[str] = None
+    graph: pyg.data.Data, title: str | None = None
 ) -> tuple[matplotlib.figure.Figure, matplotlib.axes.Axes]:
     """
     Render a PyTorch Geometric graph using stored node coordinates.
@@ -266,10 +265,10 @@ def prepend_node_index(graph: networkx.Graph, new_index: int) -> networkx.Graph:
 def create_graph(
     graph_dir_path: str,
     xy: np.ndarray,
-    n_max_levels: Optional[int] = None,
-    hierarchical: Optional[bool] = False,
-    create_plot: Optional[bool] = False,
-):
+    n_max_levels: int | None = None,
+    hierarchical: bool = False,
+    create_plot: bool = False,
+) -> None:
     """
     Create graph components from `xy` grid coordinates and store in
     `graph_dir_path`.
@@ -488,8 +487,10 @@ def create_graph(
                 .reshape((n, n, 2))[1::nx, 1::nx, :]
                 .reshape(int(n / nx) ** 2, 2)
             )
-            ij = [tuple(x) for x in ij]
-            G[lev] = networkx.relabel_nodes(G[lev], dict(zip(G[lev].nodes, ij)))
+            ij_tuples = [tuple(x) for x in ij]
+            G[lev] = networkx.relabel_nodes(
+                G[lev], dict(zip(G[lev].nodes, ij_tuples))
+            )
             G_tot = networkx.compose(G_tot, G[lev])
 
         # Relabel mesh nodes to start with 0
@@ -647,10 +648,10 @@ def create_graph(
 def create_graph_from_datastore(
     datastore: BaseRegularGridDatastore,
     output_root_path: str,
-    n_max_levels: Optional[int] = None,
+    n_max_levels: int | None = None,
     hierarchical: bool = False,
     create_plot: bool = False,
-):
+) -> None:
     """
     Generate graph components for ``datastore`` and persist them on disk.
 
@@ -683,7 +684,7 @@ def create_graph_from_datastore(
     )
 
 
-def cli(input_args: Optional[list[str]] = None) -> None:
+def cli(input_args: list[str] | None = None) -> None:
     """
     Parse CLI arguments and call :func:`create_graph_from_datastore`.
 

neural_lam/custom_loggers.py

@@ -2,7 +2,6 @@
 
 # Standard library
 import os
-from typing import Optional
 
 # Third-party
 import matplotlib.pyplot as plt
@@ -74,7 +73,7 @@ def log_image(
         self,
         key: str,
         images: list[plt.Figure],
-        step: Optional[int] = None,
+        step: int | None = None,
     ) -> None:
         """
         Log one or more matplotlib figures as images to MLFlow.

neural_lam/datastore/__init__.py

@@ -1,5 +1,8 @@
 """Datastore backends for loading and serving weather model data."""
 
+# Standard library
+from pathlib import Path
+
 # Local
 from .base import BaseDatastore  # noqa
 from .mdp import MDPDatastore  # noqa
@@ -16,7 +19,9 @@
 }
 
 
-def init_datastore(datastore_kind, config_path):
+def init_datastore(
+    datastore_kind: str, config_path: str | Path
+) -> BaseDatastore:
     """
     Instantiate a datastore based on its short-name identifier.
 

neural_lam/datastore/base.py

@@ -111,7 +111,7 @@ def get_vars_units(self, category: str) -> list[str]:
 
         Returns
         -------
-        List[str]
+        list[str]
             The units of the variables.
 
         """
@@ -127,7 +127,7 @@ def get_vars_names(self, category: str) -> list[str]:
 
         Returns
         -------
-        List[str]
+        list[str]
             The names of the variables.
 
         """
@@ -143,7 +143,7 @@ def get_vars_long_names(self, category: str) -> list[str]:
 
         Returns
         -------
-        List[str]
+        list[str]
             The long names of the variables.
 
         """
@@ -333,7 +333,7 @@ def get_xy_extent(self, category: str) -> list[float]:
 
         Returns
         -------
-        List[float]
+        list[float]
             The extent of the x, y coordinates.
 
         """
@@ -423,7 +423,7 @@ def expected_dim_order(
             The category of the dataset (state/forcing/static).
         Returns
         -------
-        List[str]
+        list[str]
             The expected dimension order for the dataarray or dataset.
 
         """
@@ -496,7 +496,7 @@ class BaseRegularGridDatastore(BaseDatastore):
     `stack_grid_coords` and `unstack_grid_coords` respectively).
     """
 
-    spatial_coordinates = ("x", "y")
+    spatial_coordinates: tuple[str, str] = ("x", "y")
 
     @cached_property
     @abc.abstractmethod
@@ -562,7 +562,7 @@ def unstack_grid_coords(
         da_or_ds_unstacked = da_or_ds.unstack("grid_index")
 
         # Ensure that the x, y dimensions are in the correct order
-        dims = da_or_ds_unstacked.dims
+        dims = list(da_or_ds_unstacked.dims)
         xy_dim_order = [d for d in dims if d in self.spatial_coordinates]
 
         if xy_dim_order != self.spatial_coordinates:
@@ -615,7 +615,7 @@ def stack_grid_coords(
         # dimension named in the format `{category}_feature`
         category = None
         for dim in da_or_ds_stacked.dims:
-            if dim.endswith("_feature"):
+            if isinstance(dim, str) and dim.endswith("_feature"):
                 if category is not None:
                     raise ValueError(
                         "Multiple dimensions ending with '_feature' found in "

neural_lam/datastore/mdp.py

@@ -7,7 +7,6 @@
 from datetime import timedelta
 from functools import cached_property
 from pathlib import Path
-from typing import List, Optional, Union
 
 # Third-party
 import cartopy.crs as ccrs
@@ -36,7 +35,7 @@ class MDPDatastore(BaseRegularGridDatastore):
 
     def __init__(
         self,
-        config_path: Union[str, Path],
+        config_path: str | Path,
         n_boundary_points: int = 30,
         reuse_existing: bool = True,
     ) -> None:
@@ -176,7 +175,7 @@ def step_length(self) -> timedelta:
         total_sec = da_dt.dt.total_seconds().isel(time=0).astype(int)
         return timedelta(seconds=int(total_sec.item()))
 
-    def get_vars_units(self, category: str) -> List[str]:
+    def get_vars_units(self, category: str) -> list[str]:
         """Return the units of the variables in the given category.
 
         Parameters
@@ -186,7 +185,7 @@ def get_vars_units(self, category: str) -> List[str]:
 
         Returns
         -------
-        List[str]
+        list[str]
             The units of the variables in the given category.
 
         """
@@ -195,7 +194,7 @@ def get_vars_units(self, category: str) -> List[str]:
             return []
         return self._ds[f"{category}_feature_units"].values.tolist()
 
-    def get_vars_names(self, category: str) -> List[str]:
+    def get_vars_names(self, category: str) -> list[str]:
         """Return the names of the variables in the given category.
 
         Parameters
@@ -205,7 +204,7 @@ def get_vars_names(self, category: str) -> List[str]:
 
         Returns
         -------
-        List[str]
+        list[str]
             The names of the variables in the given category.
 
         """
@@ -214,7 +213,7 @@ def get_vars_names(self, category: str) -> List[str]:
             return []
         return self._ds[f"{category}_feature"].values.tolist()
 
-    def get_vars_long_names(self, category: str) -> List[str]:
+    def get_vars_long_names(self, category: str) -> list[str]:
         """
         Return the long names of the variables in the given category.
 
@@ -225,7 +224,7 @@ def get_vars_long_names(self, category: str) -> List[str]:
 
         Returns
         -------
-        List[str]
+        list[str]
             The long names of the variables in the given category.
 
         """
@@ -253,9 +252,9 @@ def get_num_data_vars(self, category: str) -> int:
     def get_dataarray(
         self,
         category: str,
-        split: Optional[str],
+        split: str | None,
         standardize: bool = False,
-    ) -> Union[xr.DataArray, None]:
+    ) -> xr.DataArray | None:
         """
         Return the processed data (as a single `xr.DataArray`) for the given
         category of data and test/train/val-split that covers all the data (in
@@ -364,7 +363,7 @@ def get_standardization_dataarray(self, category: str) -> xr.Dataset:
         # Add standardized state diff stats
         if category == "state":
             ds_stats = ds_stats.assign(
-                **{
+                {
                     f"state_diff_{op}_standardized": self._ds[
                         f"state__{split}__diff_{op}"
                     ]
@@ -404,7 +403,11 @@ def boundary_mask(self) -> xr.DataArray:
         ds_unstacked["boundary_mask"] = ds_unstacked.boundary_mask.fillna(
             1
         ).astype(int)
-        return self.stack_grid_coords(da_or_ds=ds_unstacked.boundary_mask)
+        boundary_mask = self.stack_grid_coords(
+            da_or_ds=ds_unstacked.boundary_mask
+        )
+        assert isinstance(boundary_mask, xr.DataArray)
+        return boundary_mask
 
     @property
     def coords_projection(self) -> ccrs.Projection: