From 7e3cd9df637d856bbc43e8d497678304f240e894 Mon Sep 17 00:00:00 2001 From: Maarten Pronk Date: Mon, 7 Oct 2024 21:27:18 +0200 Subject: [PATCH] Allow Polygon geometry for Basin / Area. (#1875) Fixes #1850 --- python/ribasim/ribasim/geometry/area.py | 8 +++++++- python/ribasim/tests/test_validation.py | 17 ++++++++++++++++- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/python/ribasim/ribasim/geometry/area.py b/python/ribasim/ribasim/geometry/area.py index 888dc0ff2..fd9ad3e8e 100644 --- a/python/ribasim/ribasim/geometry/area.py +++ b/python/ribasim/ribasim/geometry/area.py @@ -2,7 +2,7 @@ from pandera.dtypes import Int32 from pandera.typing import Index, Series from pandera.typing.geopandas import GeoSeries -from shapely.geometry import MultiPolygon +from shapely.geometry import MultiPolygon, Polygon from .base import _GeoBaseSchema @@ -11,3 +11,9 @@ 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[MultiPolygon] = pa.Field(default=None, nullable=True) + + @pa.parser("geometry") + def convert_to_multi(cls, series): + return series.apply( + lambda geom: MultiPolygon([geom]) if isinstance(geom, Polygon) else geom + ) diff --git a/python/ribasim/tests/test_validation.py b/python/ribasim/tests/test_validation.py index cc5e8906a..28f0f6324 100644 --- a/python/ribasim/tests/test_validation.py +++ b/python/ribasim/tests/test_validation.py @@ -11,7 +11,7 @@ pid_control, pump, ) -from shapely import Point +from shapely import MultiPolygon, Point, Polygon def test_duplicate_edge(basic): @@ -171,3 +171,18 @@ def test_minimum_control_neighbor(): ], ) model.write("test.toml") + + +def test_geometry_validation(): + point = Point(0.0, 0.0) + poly = point.buffer(1.0) + + assert isinstance(poly, Polygon) + basinarea = basin.Area(geometry=[poly]) + assert isinstance(basinarea.df.geometry[0], MultiPolygon) + + basinarea = basin.Area(geometry=[basinarea.df.geometry[0]]) + assert isinstance(basinarea.df.geometry[0], MultiPolygon) + + with pytest.raises(ValueError): + basin.Area(geometry=[point])