Validate geometry types in Python.

python/ribasim/ribasim/geometry/area.py

@@ -1,14 +1,13 @@
-from typing import Any
-
 import pandera as pa
 from pandera.dtypes import Int32
 from pandera.typing import Index, Series
 from pandera.typing.geopandas import GeoSeries
+from shapely.geometry import MultiPolygon
 
-from ribasim.schemas import _BaseSchema
+from .base import _GeoBaseSchema
 
 
-class BasinAreaSchema(_BaseSchema):
+class BasinAreaSchema(_GeoBaseSchema):
     fid: Index[Int32] = pa.Field(default=0, check_name=True)
     node_id: Series[Int32] = pa.Field(nullable=False, default=0)
-    geometry: GeoSeries[Any] = pa.Field(default=None, nullable=True)
+    geometry: GeoSeries[MultiPolygon] = pa.Field(default=None, nullable=True)

python/ribasim/ribasim/geometry/base.py

@@ -0,0 +1,14 @@
+from typing import Any, get_type_hints
+
+import pandera as pa
+from pandera.typing import Series
+from pandera.typing.geopandas import GeoSeries
+
+from ribasim.schemas import _BaseSchema
+
+
+class _GeoBaseSchema(_BaseSchema):
+    @pa.check("geometry")
+    def is_correct_geometry_type(cls, geoseries: GeoSeries[Any]) -> Series[bool]:
+        T = get_type_hints(cls)["geometry"].__args__[0]
+        return geoseries.map(lambda geom: isinstance(geom, T))

python/ribasim/ribasim/geometry/edge.py

@@ -1,4 +1,4 @@
-from typing import Any, NamedTuple, Optional
+from typing import NamedTuple, Optional
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -14,9 +14,10 @@
 from shapely.geometry import LineString, MultiLineString, Point
 
 from ribasim.input_base import SpatialTableModel
-from ribasim.schemas import _BaseSchema
 from ribasim.utils import UsedIDs
 
+from .base import _GeoBaseSchema
+
 __all__ = ("EdgeTable",)
 
 SPATIALCONTROLNODETYPES = {
@@ -34,14 +35,14 @@ class NodeData(NamedTuple):
     geometry: Point
 
 
-class EdgeSchema(_BaseSchema):
+class EdgeSchema(_GeoBaseSchema):
     edge_id: Index[Int32] = pa.Field(default=0, ge=0, check_name=True)
     name: Series[str] = pa.Field(default="")
     from_node_id: Series[Int32] = pa.Field(default=0)
     to_node_id: Series[Int32] = pa.Field(default=0)
     edge_type: Series[str] = pa.Field(default="flow")
     subnetwork_id: Series[pd.Int32Dtype] = pa.Field(default=pd.NA, nullable=True)
-    geometry: GeoSeries[Any] = pa.Field(default=None, nullable=True)
+    geometry: GeoSeries[LineString] = pa.Field(default=None, nullable=True)
 
     @classmethod
     def _index_name(self) -> str:

python/ribasim/ribasim/geometry/node.py

@@ -9,21 +9,23 @@
 from pandera.dtypes import Int32
 from pandera.typing import Index, Series
 from pandera.typing.geopandas import GeoSeries
+from shapely.geometry import Point
 
 from ribasim.input_base import SpatialTableModel
-from ribasim.schemas import _BaseSchema
+
+from .base import _GeoBaseSchema
 
 __all__ = ("NodeTable",)
 
 
-class NodeSchema(_BaseSchema):
+class NodeSchema(_GeoBaseSchema):
     node_id: Index[Int32] = pa.Field(default=0, check_name=True)
     name: Series[str] = pa.Field(default="")
     node_type: Series[str] = pa.Field(default="")
     subnetwork_id: Series[pd.Int32Dtype] = pa.Field(
         default=pd.NA, nullable=True, coerce=True
     )
-    geometry: GeoSeries[Any] = pa.Field(default=None, nullable=True)
+    geometry: GeoSeries[Point] = pa.Field(default=None, nullable=True)
 
     @classmethod
     def _index_name(self) -> str:

python/ribasim/tests/test_model.py

@@ -80,7 +80,7 @@ def test_tabulated_rating_curve_model(tabulated_rating_curve, tmp_path):
     model_orig.set_crs(model_orig.crs)
     basin_area = tabulated_rating_curve.basin.area.df
     assert basin_area is not None
-    assert basin_area.geometry.geom_type.iloc[0] == "Polygon"
+    assert basin_area.geometry.geom_type.iloc[0] == "MultiPolygon"
     assert basin_area.crs == CRS.from_epsg(28992)
     model_orig.write(tmp_path / "tabulated_rating_curve/ribasim.toml")
     model_new = Model.read(tmp_path / "tabulated_rating_curve/ribasim.toml")

python/ribasim/tests/test_schemas.py

@@ -1,6 +1,18 @@
+import pytest
+from pydantic import ValidationError
+from ribasim.nodes import basin
 from ribasim.schemas import BasinProfileSchema
+from shapely.geometry import Point
 
 
 def test_config_inheritance():
     assert BasinProfileSchema.__config__.add_missing_columns
     assert BasinProfileSchema.__config__.coerce
+
+
+def test_geometry_validation():
+    with pytest.raises(
+        ValidationError,
+        match="Column 'geometry' failed element-wise validator number 0: <Check is_correct_geometry_type> failure cases",
+    ):
+        basin.Area(geometry=[Point([1.0, 2.0])])

python/ribasim_testmodels/ribasim_testmodels/basic.py

@@ -17,7 +17,7 @@
     pump,
     tabulated_rating_curve,
 )
-from shapely.geometry import Point
+from shapely.geometry import MultiPolygon, Point
 
 
 def basic_model() -> ribasim.Model:
@@ -314,7 +314,7 @@ def tabulated_rating_curve_model() -> ribasim.Model:
         Node(1, basin_geometry_1),
         [
             basin.Static(precipitation=[0.002 / 86400]),
-            basin.Area(geometry=[basin_geometry_1.buffer(1.0)]),
+            basin.Area(geometry=[MultiPolygon([basin_geometry_1.buffer(1.0)])]),
             *node_data,
         ],
     )
@@ -323,7 +323,7 @@ def tabulated_rating_curve_model() -> ribasim.Model:
         Node(4, basin_geometry_2),
         [
             basin.Static(precipitation=[0.0]),
-            basin.Area(geometry=[basin_geometry_2.buffer(1.0)]),
+            basin.Area(geometry=[MultiPolygon([basin_geometry_2.buffer(1.0)])]),
             *node_data,
         ],
     )

ribasim_qgis/core/nodes.py

@@ -373,7 +373,7 @@
 
     @classmethod
     def geometry_type(cls) -> str:
-        return "Polygon"
+        return "MultiPolygon"
 
     @classmethod
     def qgis_geometry_type(cls) -> Qgis.GeometryType:

utils/generate-testmodels.py

@@ -1,10 +1,15 @@
 import multiprocessing
 import shutil
+import sys
 from functools import partial
 from pathlib import Path
 
 import ribasim_testmodels
 
+selection = (
+    sys.argv[1:] if len(sys.argv) > 1 else ribasim_testmodels.constructors.keys()
+)
+
 
 def generate_model(args, datadir):
     model_name, model_constructor = args
@@ -30,8 +35,10 @@ def generate_model(args, datadir):
 
     generate_model_partial = partial(generate_model, datadir=datadir)
 
+    models = [
+        (k, v) for k, v in ribasim_testmodels.constructors.items() if k in selection
+    ]
+
     with multiprocessing.Pool(processes=4) as p:
-        for model_name in p.imap_unordered(
-            generate_model_partial, ribasim_testmodels.constructors.items()
-        ):
+        for model_name in p.imap_unordered(generate_model_partial, models):
             print(f"Generated {model_name}")