Better type hints

src/itzi_core/data_containers.py

@@ -19,7 +19,15 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
-from pydantic import BaseModel, ConfigDict, Field, NonNegativeFloat, NonNegativeInt, PositiveFloat
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    NonNegativeFloat,
+    NonNegativeInt,
+    PositiveFloat,
+    field_validator,
+)
 
 from itzi_core.const import DefaultValues, InfiltrationModelType, TemporalType
 from itzi_core.providers.domain_data import DomainData
@@ -146,7 +154,7 @@ class SimulationData(BaseModel):
     sim_time: datetime
     time_step: float  # time step duration
     time_steps_counter: int  # number of time steps since last update
-    continuity_data: ContinuityData | None  # Made optional for use in tests
+    continuity_data: ContinuityData
     raw_arrays: dict[str, np.ndarray]
     accumulation_arrays: dict[str, np.ndarray]
     cell_dx: PositiveFloat  # cell size in east-west direction
@@ -211,8 +219,8 @@ class SimulationConfig(BaseModel):
     # Hotstart config
     hotstart_config: HotstartRunConfig | None = None
     # Input and output raster maps
-    input_map_names: dict[str, str | None]
-    output_map_names: dict[str, str | None]
+    input_map_names: dict[str, str]
+    output_map_names: dict[str, str]
     # Surface flow parameters
     surface_flow_parameters: SurfaceFlowParameters
     # Hydrology parameters
@@ -225,6 +233,14 @@ class SimulationConfig(BaseModel):
     free_weir_coeff: NonNegativeFloat = Field(DefaultValues.FREE_WEIR_COEFF, ge=0, le=1)
     submerged_weir_coeff: NonNegativeFloat = Field(DefaultValues.SUBMERGED_WEIR_COEFF, ge=0, le=1)
 
+    @field_validator("input_map_names", "output_map_names", mode="before")
+    @classmethod
+    def remove_inactive_map_names(cls, value: object) -> object:
+        """Normalize legacy null-valued map entries to omitted inactive entries."""
+        if isinstance(value, dict):
+            return {key: map_name for key, map_name in value.items() if map_name is not None}
+        return value
+
     def as_str_dict(self) -> dict:
         """Convert the configuration to a dictionary with string representations."""
         raw_dict = self.model_dump()

src/itzi_core/drainage.py

@@ -14,27 +14,29 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
 import math
+import tempfile
 from datetime import timedelta
 from enum import StrEnum
 from io import BytesIO
-import tempfile
+from typing import TYPE_CHECKING
 
-import pyswmm
 import numpy as np
+import pyswmm
+from pyswmm.toolkitapi import NodeResults, SimulationParameters, SimulationTime
 
 from itzi_core import DefaultValues
 from itzi_core.data_containers import (
-    DrainageNodeData,
+    DrainageLinkAttributes,
     DrainageLinkData,
     DrainageNetworkData,
-    DrainageLinkAttributes,
     DrainageNodeAttributes,
+    DrainageNodeData,
 )
 
 if TYPE_CHECKING:
     from datetime import datetime
+
     from pyswmm.swmm5 import PySWMM
 
 
@@ -83,7 +85,7 @@ def __init__(
             self.swmm_model.swmm_use_hotstart(hotstart_filename)
         if hotstart_start_datetime is not None:
             self.swmm_model.setSimulationDateTime(
-                pyswmm.toolkitapi.SimulationTime.StartDateTime, hotstart_start_datetime
+                SimulationTime.StartDateTime, hotstart_start_datetime
             )
         self.swmm_model.swmm_start()
         # allow ponding
@@ -200,9 +202,7 @@ def __init__(
         self.free_weir_coeff = free_weir_coeff
         self.submerged_weir_coeff = submerged_weir_coeff
         self.node_type = self.get_node_type()
-        self.surface_area = self._model.getSimAnalysisSetting(
-            pyswmm.toolkitapi.SimulationParameters.MinSurfArea
-        )
+        self.surface_area = self._model.getSimAnalysisSetting(SimulationParameters.MinSurfArea)
         # weir width is the circumference (node considered circular)
         self.weir_width = 2 * math.sqrt(self.surface_area * math.pi)
         # Set default values
@@ -229,7 +229,7 @@ def get_full_volume(self):
         return self.surface_area * self.pyswmm_node.full_depth
 
     def get_overflow(self):
-        return self._model.getNodeResult(self.node_id, pyswmm.toolkitapi.NodeResults.overflow)
+        return self._model.getNodeResult(self.node_id, NodeResults.overflow)
 
     def get_crest_elev(self):
         """Return the crest elevation of the node."""

src/itzi_core/itzi_error.py

@@ -16,34 +16,14 @@
 class NullError(RuntimeError):
     """Raised when null values is detected in simulation"""
 
-    pass
-
 
 class DtError(RuntimeError):
     """Error related to time-step calculation"""
 
-    def __init__(self, msg):
-        self.msg = msg
-
-    def __str__(self):
-        return repr(self.msg)
-
 
 class MassBalanceError(RuntimeError):
     """Raised when mass balance error exceeds threshold"""
 
-    def __init__(self, msg: str):
-        self.msg = msg
-
-    def __str__(self):
-        return repr(self.msg)
-
 
 class HotstartError(RuntimeError):
     """Raised when hotstart file operations fail."""
-
-    def __init__(self, msg):
-        self.msg = msg
-
-    def __str__(self):
-        return repr(self.msg)

src/itzi_core/profiler.py

@@ -18,8 +18,9 @@
 """
 
 import os
-from pathlib import Path
+from collections.abc import Iterator
 from contextlib import contextmanager
+from pathlib import Path
 
 # Attempt to import pyinstrument
 try:
@@ -32,7 +33,7 @@
 
 
 @contextmanager
-def profile_context(file_path: Path = None):
+def profile_context(file_path: Path | None = None) -> Iterator[None]:
     """
     A context manager for profiling code blocks.
 
@@ -49,7 +50,9 @@ def profile_context(file_path: Path = None):
             # Code to be profiled (or not)
             run_simulation()
     """
-    profiler_active = os.environ.get("ITZI_PROFILE") == "1" and PYINSTRUMENT_AVAILABLE
+    profiler_active = (
+        os.environ.get("ITZI_PROFILE") == "1" and PYINSTRUMENT_AVAILABLE and Profiler is not None
+    )
 
     if profiler_active:
         profiler = Profiler()

src/itzi_core/providers/csv_output.py

@@ -18,7 +18,7 @@
 from datetime import datetime, timedelta
 from io import StringIO
 from pathlib import PurePosixPath, PureWindowsPath
-from typing import TYPE_CHECKING, TypedDict
+from typing import TYPE_CHECKING, Any, TypedDict
 
 import pandas as pd
 
@@ -61,10 +61,8 @@ class CSVVectorOutputProvider(VectorOutputProvider):
 
     def __init__(self, config: CSVVectorOutputConfig) -> None:
         """Initialize output provider with provider configuration."""
-        try:
-            self.srid = config["crs"].to_epsg()
-        except AttributeError:
-            self.srid = 0
+        crs = config["crs"]
+        self.srid = 0 if crs is None else crs.to_epsg() or 0
         self.store = config["store"]
         prefix_str = config["results_prefix"]
         prefix_path = PurePosixPath(prefix_str.replace("\\", "/"))
@@ -76,11 +74,14 @@ def __init__(self, config: CSVVectorOutputConfig) -> None:
             raise ValueError("results_prefix must be a relative path without parent traversal")
         results_prefix = prefix_path.as_posix() if prefix_path.parts else ""
 
-        self.existing_ids = {"link": None, "node": None}  # Objects ids already in the file
-        self.existing_max_time = {"link": None, "node": None}  # Max of sim_time in existing_file
+        self.existing_ids: dict[str, set[Any] | None] = {"link": None, "node": None}
+        self.existing_max_time: dict[str, datetime | timedelta | None] = {
+            "link": None,
+            "node": None,
+        }
         self.number_of_writes = {"link": 0, "node": 0}
-        self.file_paths = {"link": None, "node": None}
-        self.headers = {"link": None, "node": None}
+        self.file_paths: dict[str, str] = {}
+        self.headers: dict[str, list[str]] = {}
         self.append_mode = {"link": True, "node": True}
         if config["overwrite"]:
             self.append_mode = {"link": False, "node": False}
@@ -97,8 +98,6 @@ def __init__(self, config: CSVVectorOutputConfig) -> None:
             # create the CSV files
             if not self.append_mode[geom_type]:
                 self._write_headers(geom_type)
-        print(self.existing_ids)
-        print(self.existing_max_time)
 
     def write_vector(
         self, drainage_data: DrainageNetworkData, sim_time: datetime | timedelta
@@ -130,7 +129,6 @@ def _check_existing_csv(self, geom_type: str):
             - new object ID ≠ existing ones
         could not be checked without drainage network data
         """
-        existing_csv = None
         try:
             existing_csv = StringIO(
                 bytes(obstore.get(self.store, self.file_paths[geom_type]).bytes()).decode("utf-8")
@@ -143,7 +141,6 @@ def _check_existing_csv(self, geom_type: str):
         expected_headers = self.headers[geom_type]
         if not existing_headers == expected_headers:
             raise ValueError(f"Headers mismatch in existing file {self.file_paths[geom_type]}.")
-            self.append_mode[geom_type] = False
         id_col = f"{geom_type}_id"
 
         # Store values existing ids
@@ -162,7 +159,6 @@ def _check_existing_csv(self, geom_type: str):
                 raise ValueError(
                     f"Unknown sim_time column in existing file {self.file_paths[geom_type]}."
                 )
-        print(df_csv)
 
     def _write_headers(self, geom_type: str):
         """Create an in-memory CSV file with headers and save it in the store."""
@@ -178,29 +174,33 @@ def _validate_time_on_first_write(self, sim_time: datetime | timedelta) -> None:
 
         for geom_type in ["node", "link"]:
             # Only validate on first write
-            if self.number_of_writes[geom_type] > 0 or self.existing_max_time[geom_type] is None:
+            existing_time = self.existing_max_time[geom_type]
+            if self.number_of_writes[geom_type] > 0 or existing_time is None:
                 continue
-            # Type must match
-            if type(self.existing_max_time[geom_type]) is not type(sim_time):
+            # ty finds an error when only one if is used
+            # ruff: noqa: SIM114
+            if isinstance(sim_time, datetime) and isinstance(existing_time, datetime):
+                time_is_increasing = sim_time > existing_time
+            elif isinstance(sim_time, timedelta) and isinstance(existing_time, timedelta):
+                time_is_increasing = sim_time > existing_time
+            else:
                 time_type_name = (
                     "relative (timedelta)"
                     if isinstance(sim_time, timedelta)
                     else "absolute (datetime)"
                 )
                 existing_type_name = (
-                    "relative"
-                    if isinstance(self.existing_max_time[geom_type], timedelta)
-                    else "absolute"
+                    "relative" if isinstance(existing_time, timedelta) else "absolute"
                 )
-                raise ValueError(
+                raise TypeError(
                     f"Time type mismatch for {geom_type}: "
                     f"attempting to write {time_type_name} but existing file has {existing_type_name}"
                 )
             # Time must increase
-            if not sim_time > self.existing_max_time[geom_type]:
+            if not time_is_increasing:
                 raise ValueError(
                     f"Time not increasing for {geom_type}: attempting to write {sim_time} but "
-                    f"existing file has a max sim_time value of {self.existing_max_time[geom_type]}"
+                    f"existing file has a max sim_time value of {existing_time}"
                 )
 
     def _update_csv(
@@ -235,7 +235,7 @@ def _update_csv(
         updated_csv = existing_csv + new_rows
         obstore.put(self.store, self.file_paths[geom_type], file=updated_csv.encode("utf-8"))
 
-    def _attrs_line(self, drainage_element: DrainageNodeData | DrainageLinkData) -> list[str, ...]:
+    def _attrs_line(self, drainage_element: DrainageNodeData | DrainageLinkData) -> list[str]:
         """Return a list of attributes"""
         # Convert attributes to list
         attributes = [str(a) for a in drainage_element.attributes.model_dump().values()]

src/itzi_core/providers/memory_output.py

@@ -29,7 +29,9 @@ def __init__(self, out_map_names: Mapping[str, str]) -> None:
         """Initialize output provider with simulation configuration."""
         # user-selected map names.
         self.out_map_names = out_map_names
-        self.output_maps_dict: dict[str, list] = {k: [] for k in self.out_map_names}
+        self.output_maps_dict: dict[str, list[tuple[datetime | timedelta, np.ndarray]]] = {
+            key: [] for key in self.out_map_names
+        }
 
     def write_arrays(
         self, array_dict: Mapping[str, np.ndarray], sim_time: datetime | timedelta

src/itzi_core/providers/xarray_input.py

@@ -14,23 +14,24 @@
 
 from __future__ import annotations
 
-from typing import Iterable, Mapping, TypedDict, NotRequired, TYPE_CHECKING
+from collections.abc import Iterable, Mapping
+from typing import TYPE_CHECKING, NotRequired, TypedDict
 
 import numpy as np
 
 try:
-    import xarray as xr
     import pandas as pd
+    import xarray as xr
 except ImportError:
     raise ImportError(
         "To use the xarray input backend, install itzi with: "
         "'uv tool install itzi[cloud]' "
         "or 'pip install itzi[cloud]'"
     )
 
+from itzi_core.const import TemporalType
 from itzi_core.providers.base import RasterInputProvider
 from itzi_core.providers.domain_data import DomainData
-from itzi_core.const import TemporalType
 
 if TYPE_CHECKING:
     from datetime import datetime
@@ -148,9 +149,10 @@ def _validate_dimensions(self) -> None:
         """Validate that:
         - the specified spatial dimensions exist in the dataset.
         - The dimensions are one-dimensional."""
-        for var_name in self.dataset.data_vars.keys():
+        for raw_var_name in self.dataset.data_vars.keys():
+            var_name = str(raw_var_name)
             da_var: xr.DataArray = self.dataset[var_name]
-            var_dims: set[str] = set(da_var.dims)
+            var_dims = {str(dim) for dim in da_var.dims}
             for dim_type in ["x", "y"]:
                 dim_name: str = self.dataset_dims[var_name][dim_type]
                 if dim_name not in var_dims:
@@ -169,7 +171,7 @@ def _validate_variables_dimensionality(self) -> None:
         """Validate that all variables are either 2D[y, x] or 3D[time, y, x]."""
         for var_name in self.input_map_names.values():
             da_var: xr.DataArray = self.dataset[var_name]
-            var_dims: set[str] = set(da_var.dims)
+            var_dims = {str(dim) for dim in da_var.dims}
             num_dims = len(da_var.dims)
             x_dim: str = self.dataset_dims[var_name]["x"]
             y_dim: str = self.dataset_dims[var_name]["y"]
@@ -213,7 +215,7 @@ def _validate_equal_spacing_of_spatial_dims(self) -> None:
                 diffs = np.diff(coord.values if hasattr(coord, "values") else coord)
                 # no coordinates present
                 if len(diffs) == 0:
-                    pass
+                    continue
                 if not np.allclose(diffs, diffs[0]):
                     raise ValueError(
                         f"Dimension {dim_name} of variable {var_name} not equally spaced."
@@ -231,7 +233,7 @@ def _validate_equality_of_spatial_dims(self):
             if len(dim_names) == 0:
                 continue
             da_list: list[xr.DataArray] = [self.dataset[dim_name] for dim_name in dim_names]
-            ref_da: np.ndarray = da_list[0]
+            ref_da: xr.DataArray = da_list[0]
             for da in da_list:
                 if not np.allclose(ref_da.values, da.values):
                     raise ValueError(