Replace dataclasses by pydantic BaseModel (#209)

* data_container.py transition from dataclasses to pydantic BaseModel

* remove fail-fast from build wheel

* build swmm-toolkit without stable ABI

---------

Co-authored-by: karosc <[email protected]>

Author: lrntct

.github/workflows/build_wheels.yml

@@ -10,6 +10,7 @@ jobs:
     name: Build wheels on ${{ matrix.os }}
     runs-on: ${{ matrix.os }}
     strategy:
+      fail-fast: false
       matrix:
         # macos-13 is intel, macos-14 is apple silicon, windows-11-arm
         os: [ubuntu-latest, ubuntu-24.04-arm, windows-latest, windows-11-arm]
@@ -27,10 +28,11 @@ jobs:
         env:
           CIBW_ARCHS: auto64  # Don't build 32 bits wheels
           CIBW_SKIP: pp*  # Don't build pypy wheels
-          CIBW_ENVIRONMENT: "ITZI_BDIST_WHEEL=1"  # Trigger generic compiler flags
+          CIBW_ENVIRONMENT: "ITZI_BDIST_WHEEL=1 NO_STABLE_ABI=1"  # Trigger generic compiler flags
           CIBW_ENVIRONMENT_MACOS: >
             MACOSX_DEPLOYMENT_TARGET=14.0
-            ITZI_BDIST_WHEEL=1
+            ITZI_BDIST_WHEEL=1 
+            NO_STABLE_ABI=1
           CIBW_BEFORE_TEST: "pip install pytest pytest-benchmark pandas"
           CIBW_TEST_SOURCES: tests
           CIBW_TEST_COMMAND: >  # Fast tests which do not require GRASS

pyproject.toml

@@ -9,6 +9,7 @@ dependencies = [
     "numpy>=2.2",
     "pyswmm>=2.1.0",
     "bmipy>=2.0.1",
+    "pydantic>=2.12.5",
 ]
 requires-python = ">=3.11,<3.14"  # pyswmm does not support python > 3.13
 readme = "README.md"

src/itzi/data_containers.py

@@ -12,22 +12,26 @@
 GNU General Public License for more details.
 """
 
+from __future__ import annotations
+
 from typing import Dict, Tuple, TYPE_CHECKING
-import dataclasses
 from datetime import datetime, timedelta
+from pathlib import Path
 
 import numpy as np
+from pydantic import BaseModel, ConfigDict
 
 from itzi.const import DefaultValues, TemporalType, InfiltrationModelType
 
 if TYPE_CHECKING:
     from itzi.drainage import DrainageNode
 
 
-@dataclasses.dataclass(frozen=True)
-class DrainageNodeCouplingData:
+class DrainageNodeCouplingData(BaseModel):
     """Store the translation between coordinates and array location for a given drainage node."""
 
+    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
+
     node_id: str  # Name of the drainage node
     node_object: "DrainageNode"
     # Location in the coordinate system
@@ -38,24 +42,24 @@ class DrainageNodeCouplingData:
     col: int | None
 
 
-@dataclasses.dataclass(frozen=True)
-class DrainageAttributes:
+class DrainageAttributes(BaseModel):
     """A base class for drainage data attributes."""
 
+    model_config = ConfigDict(frozen=True)
+
     @classmethod
     def get_columns_definition(cls, cat_primary_key=True) -> list[tuple[str, str]]:
         """Return a list of tuples to create DB columns"""
         type_mapping = {str: "TEXT", int: "INT", float: "REAL"}
         db_columns_def = [("cat", "INTEGER PRIMARY KEY")]
         if not cat_primary_key:
             db_columns_def = []
-        for f in dataclasses.fields(cls):
-            db_field = (f.name, type_mapping[f.type])
+        for field_name, field_info in cls.model_fields.items():
+            db_field = (field_name, type_mapping[field_info.annotation])
             db_columns_def.append(db_field)
         return db_columns_def
 
 
-@dataclasses.dataclass(frozen=True)
 class DrainageLinkAttributes(DrainageAttributes):
     link_id: str
     link_type: str
@@ -67,16 +71,16 @@ class DrainageLinkAttributes(DrainageAttributes):
     froude: float
 
 
-@dataclasses.dataclass(frozen=True)
-class DrainageLinkData:
+class DrainageLinkData(BaseModel):
     """Store the instantaneous state of a node during a drainage simulation.
     Vertices include the coordinates of the start and end nodes."""
 
-    vertices: None | Tuple[Tuple[float, float], ...]
+    model_config = ConfigDict(frozen=True)
+
+    vertices: None | Tuple[Tuple[float, float] | None, ...]
     attributes: DrainageLinkAttributes
 
 
-@dataclasses.dataclass(frozen=True)
 class DrainageNodeAttributes(DrainageAttributes):
     node_id: str
     node_type: str
@@ -101,55 +105,60 @@ class DrainageNodeAttributes(DrainageAttributes):
     full_volume: float
 
 
-@dataclasses.dataclass(frozen=True)
-class DrainageNodeData:
+class DrainageNodeData(BaseModel):
     """Store the instantaneous state of a node during a drainage simulation"""
 
+    model_config = ConfigDict(frozen=True)
+
     coordinates: None | Tuple[float, float]
     attributes: DrainageNodeAttributes
 
 
-@dataclasses.dataclass(frozen=True)
-class DrainageNetworkData:
+class DrainageNetworkData(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
     nodes: Tuple[DrainageNodeData, ...]
     links: Tuple[DrainageLinkData, ...]
 
 
-@dataclasses.dataclass(frozen=True)
-class ContinuityData:
+class ContinuityData(BaseModel):
     """Store information about simulation continuity"""
 
+    model_config = ConfigDict(frozen=True)
+
     new_domain_vol: float
     volume_change: float
     volume_error: float
     continuity_error: float
 
 
-@dataclasses.dataclass(frozen=True)
-class SimulationData:
+class SimulationData(BaseModel):
     """Immutable data container for passing raw simulation state to Report.
 
     This is a pure data structure containing only the "raw ingredients"
     needed for a report. All report-specific calculations (e.g., WSE,
     average rates) are performed by the Report class itself.
     """
 
+    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
+
     sim_time: datetime
     time_step: float  # time step duration
     time_steps_counter: int  # number of time steps since last update
-    continuity_data: ContinuityData
+    continuity_data: ContinuityData | None  # Made optional for use in tests
     raw_arrays: Dict[str, np.ndarray]
     accumulation_arrays: Dict[str, np.ndarray]
     cell_dx: float  # cell size in east-west direction
     cell_dy: float  # cell size in north-south direction
     drainage_network_data: DrainageNetworkData | None
 
 
-@dataclasses.dataclass(frozen=True)
-class MassBalanceData:
+class MassBalanceData(BaseModel):
     """Contains the fields written to the mass balance file"""
 
-    simulation_time: datetime
+    model_config = ConfigDict(frozen=True)
+
+    simulation_time: datetime | timedelta
     average_timestep: float
     timesteps: int
     boundary_volume: float
@@ -164,10 +173,11 @@ class MassBalanceData:
     percent_error: float
 
 
-@dataclasses.dataclass(frozen=True)
-class SurfaceFlowParameters:
+class SurfaceFlowParameters(BaseModel):
     """Parameters for the surface flow model."""
 
+    model_config = ConfigDict(frozen=True)
+
     hmin: float = DefaultValues.HFMIN
     cfl: float = DefaultValues.CFL
     theta: float = DefaultValues.THETA
@@ -179,22 +189,23 @@ class SurfaceFlowParameters:
     max_error: float = DefaultValues.MAX_ERROR
 
 
-@dataclasses.dataclass(frozen=True)
-class SimulationConfig:
+class SimulationConfig(BaseModel):
     """Configuration data for a simulation run."""
 
+    model_config = ConfigDict(frozen=True)
+
     # Simulation times
     start_time: datetime
     end_time: datetime
     record_step: timedelta
     temporal_type: TemporalType
     # Input and output raster maps
-    input_map_names: Dict[str, str]
-    output_map_names: Dict[str, str]
+    input_map_names: Dict[str, str | None]
+    output_map_names: Dict[str, str | None]
     # Surface flow parameters
     surface_flow_parameters: SurfaceFlowParameters
     # Mass balance file
-    stats_file: str
+    stats_file: str | Path | None = None
     # Hydrology parameters
     dtinf: float = DefaultValues.DTINF
     infiltration_model: InfiltrationModelType = InfiltrationModelType.NULL

src/itzi/drainage.py

@@ -352,3 +352,9 @@ def get_attrs(self) -> DrainageLinkAttributes:
 
     def get_data(self) -> DrainageLinkData:
         return DrainageLinkData(vertices=self.vertices, attributes=self.get_attrs())
+
+
+# Rebuild Pydantic models that have forward references to DrainageNode
+from itzi.data_containers import DrainageNodeCouplingData  # noqa E402
+
+DrainageNodeCouplingData.model_rebuild()

src/itzi/massbalance.py

@@ -15,7 +15,6 @@
 from datetime import datetime
 import csv
 import numbers
-import dataclasses
 
 from itzi.data_containers import MassBalanceData
 
@@ -28,7 +27,7 @@ def __init__(
         file_name: str,
     ):
         """Initializes the logger and creates the output file with headers."""
-        self.fields = [f.name for f in dataclasses.fields(MassBalanceData)]
+        self.fields = list(MassBalanceData.model_fields.keys())
         self.file_name = self._set_file_name(file_name)
         self._create_file()
 
@@ -48,7 +47,7 @@ def log(self, report_data: MassBalanceData) -> None:
         """Writes a single line of data to the CSV file."""
         line_to_write = {}
 
-        for key, value in dataclasses.asdict(report_data).items():
+        for key, value in report_data.model_dump().items():
             if value != value:  # test for NaN
                 line_to_write[key] = "-"
             elif "percent_error" == key:

src/itzi/providers/csv_output.py

@@ -18,7 +18,6 @@
 from typing import TypedDict, TYPE_CHECKING, Tuple, List
 from io import StringIO
 import csv
-import dataclasses
 
 import pandas as pd
 
@@ -74,8 +73,8 @@ def __init__(self, config: CSVVectorOutputConfig) -> None:
             self.append_mode = {"link": False, "node": False}
 
         for geom_type, obj in [("node", DrainageNodeAttributes), ("link", DrainageLinkAttributes)]:
-            base_headers = [field.name for field in dataclasses.fields(obj)]
-            self.headers[geom_type] = ["sim_time"] + list(base_headers) + ["srid", "geometry"]
+            base_headers = list(obj.model_fields.keys())
+            self.headers[geom_type] = ["sim_time"] + base_headers + ["srid", "geometry"]
 
             results_name = f"{config['drainage_results_name']}_{geom_type}s.csv"
             self.file_paths[geom_type] = results_prefix + "/" + results_name
@@ -207,7 +206,7 @@ def _update_csv(
             # IDs must match
             new_ids = set(
                 [
-                    dataclasses.asdict(drainage_elem.attributes)[f"{geom_type}_id"]
+                    drainage_elem.attributes.model_dump()[f"{geom_type}_id"]
                     for drainage_elem in drainage_elements
                 ]
             )
@@ -232,7 +231,7 @@ def _update_csv(
     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 dataclasses.asdict(drainage_element.attributes).values()]
+        attributes = [str(a) for a in drainage_element.attributes.model_dump().values()]
         # Create geometry WKT
         if isinstance(drainage_element, DrainageNodeData):
             if drainage_element.coordinates is not None:

src/itzi/providers/grass_interface.py

@@ -16,7 +16,6 @@
 from pathlib import Path
 from collections import namedtuple
 from datetime import datetime, timedelta
-import dataclasses
 from typing import Self
 
 # from multiprocessing import Process, JoinableQueue
@@ -225,7 +224,6 @@ def finalize(self):
 
     def cleanup(self):
         """Remove temporary region and mask."""
-        msgr.debug("Reset mask and region")
         if self.raster_mask_id:
             msgr.debug("Remove temp MASK...")
             self.del_temp_mask()
@@ -542,9 +540,7 @@ def write_vector_map(self, drainage_data: DrainageNetworkData, map_name: str) ->
                     # The write function of the vector map set the layer to the one we set earlier
                     vector_map.write(point, cat=cat_num)
                 # Get DB attributes even if no associated geometry
-                node_attributes = tuple(
-                    value for _, value in dataclasses.asdict(node.attributes).items()
-                )
+                node_attributes = tuple(value for _, value in node.attributes.model_dump().items())
                 attrs = (cat_num,) + node_attributes
                 db_info["node"].append(attrs)
                 cat_num += 1
@@ -560,9 +556,7 @@ def write_vector_map(self, drainage_data: DrainageNetworkData, map_name: str) ->
                     # The write function of the vector map set the layer to the one we set earlier
                     vector_map.write(line_object, cat=cat_num)
                 # Get DB attributes even if no associated geometry
-                link_attributes = tuple(
-                    value for _, value in dataclasses.asdict(link.attributes).items()
-                )
+                link_attributes = tuple(value for _, value in link.attributes.model_dump().items())
                 attrs = (cat_num,) + link_attributes
                 db_info["link"].append(attrs)
                 cat_num += 1

src/itzi/simulation.py

@@ -413,7 +413,12 @@ def get_continuity_data(self) -> ContinuityData:
         else:
             continuity_error = volume_error / volume_change
 
-        return ContinuityData(new_domain_vol, volume_change, volume_error, continuity_error)
+        return ContinuityData(
+            new_domain_vol=new_domain_vol,
+            volume_change=volume_change,
+            volume_error=volume_error,
+            continuity_error=continuity_error,
+        )
 
     def _update_accum_array(self, k: str, sim_time: datetime) -> None:
         """Update the accumulation arrays."""

tests/grass_tests/test_itzi.py

@@ -2,7 +2,6 @@
 
 import os
 from io import StringIO
-import dataclasses
 
 import pandas as pd
 import numpy as np
@@ -123,7 +122,7 @@ def test_stats_file(test_data_temp_path):
     assert os.path.exists(stats_path)
     df = pd.read_csv(stats_path)
 
-    expected_cols = [f.name for f in dataclasses.fields(MassBalanceData)]
+    expected_cols = list(MassBalanceData.model_fields.keys())
     assert df.columns.to_list() == expected_cols
 
     # Domain area in m2

tests/grass_tests/test_tutorial.py

@@ -6,7 +6,6 @@
 import pathlib
 import tempfile
 from configparser import ConfigParser
-from dataclasses import fields
 
 import pytest
 import numpy as np
@@ -272,7 +271,7 @@ def test_tutorial_drainage(itzi_tutorial, test_data_path, test_data_temp_path, h
         assert len(link_entries) == 2
         # link DB columns are as expected
         actual_link_columns = v_db_select[0].split("|")
-        expected_link_columns = ["cat"] + [field.name for field in fields(DrainageLinkAttributes)]
+        expected_link_columns = ["cat"] + list(DrainageLinkAttributes.model_fields.keys())
         assert expected_link_columns == actual_link_columns
 
         # Check nodes DB table
@@ -282,5 +281,5 @@ def test_tutorial_drainage(itzi_tutorial, test_data_path, test_data_temp_path, h
         assert len(nodes_entries) == 3
         # node DB columns are as expected
         actual_nodes_columns = v_db_select[0].split("|")
-        expected_nodes_columns = ["cat"] + [field.name for field in fields(DrainageNodeAttributes)]
+        expected_nodes_columns = ["cat"] + list(DrainageNodeAttributes.model_fields.keys())
         assert expected_nodes_columns == actual_nodes_columns