Fix model plotting while still ensuring that Edge tables are written

python/ribasim/ribasim/geometry/edge.py

@@ -8,7 +8,7 @@
 from geopandas import GeoDataFrame
 from matplotlib.axes import Axes
 from numpy.typing import NDArray
-from pandera.typing import DataFrame, Series
+from pandera.typing import Series
 from pandera.typing.geopandas import GeoSeries
 from shapely.geometry import LineString, MultiLineString, Point
 
@@ -42,10 +42,6 @@
 class EdgeTable(SpatialTableModel[EdgeSchema]):
     """Defines the connections between nodes."""
 
-    def __init__(self, **kwargs):
-        kwargs.setdefault("df", DataFrame[EdgeSchema]())
-        super().__init__(**kwargs)
-
     def add(
         self,
         from_node: NodeData,
@@ -82,8 +78,9 @@
         assert self.df is not None
         return (self.df.edge_type == edge_type).to_numpy()
 
-    def plot(self, **kwargs) -> Axes:
-        assert self.df is not None  # Pleases mypy
+    def plot(self, **kwargs) -> Axes | None:
+        if self.df is None:
+            return None
         kwargs = kwargs.copy()  # Avoid side-effects
         ax = kwargs.get("ax", None)
         color_flow = kwargs.pop("color_flow", None)

python/ribasim/ribasim/geometry/node.py

@@ -15,7 +15,7 @@
 
 
 class NodeSchema(pa.SchemaModel):
-    node_id: Series[int]
+    node_id: Series[int] = pa.Field(nullable=False, default=0)
     name: Series[str] = pa.Field(default="")
     node_type: Series[str] = pa.Field(default="")
     subnetwork_id: Series[pd.Int64Dtype] = pa.Field(

python/ribasim/ribasim/input_base.py

@@ -235,7 +235,7 @@ def _save(
         if self.df is not None and self.filepath is not None:
             self.sort()
             self._write_arrow(self.filepath, directory, input_dir)
-        elif self.df is not None and db_path is not None:
+        elif db_path is not None:
             self.sort()
             self._write_table(db_path)
 
@@ -358,8 +358,8 @@ def _write_table(self, path: FilePath) -> None:
         ----------
         path : FilePath
         """
-
-        gdf = gpd.GeoDataFrame(data=self.df)
+        df = DataFrame[self.tableschema()]() if self.df is None else self.df  # type:ignore
+        gdf = gpd.GeoDataFrame(data=df)
         gdf = gdf.set_geometry("geometry")
         gdf.to_file(path, layer=self.tablename(), driver="GPKG", mode="a")