From 8d7af62525d419020c2246ea45a0228f57c3824b Mon Sep 17 00:00:00 2001
From: Martijn Visser <mgvisser@gmail.com>
Date: Fri, 2 Feb 2024 13:47:38 +0100
Subject: [PATCH] Add asserts to satisfy mypy

Node and Edge tables are not allowed missing, though we say they are in TableModel from which they inherit.

(cherry picked from commit a263cee52c20dba59e7cf789edb53374ede74248)
---
 python/ribasim/ribasim/geometry/edge.py | 2 ++
 python/ribasim/ribasim/geometry/node.py | 5 +++++
 python/ribasim/ribasim/model.py         | 2 ++
 3 files changed, 9 insertions(+)

diff --git a/python/ribasim/ribasim/geometry/edge.py b/python/ribasim/ribasim/geometry/edge.py
index 7148e9ee1..3abd3216e 100644
--- a/python/ribasim/ribasim/geometry/edge.py
+++ b/python/ribasim/ribasim/geometry/edge.py
@@ -40,12 +40,14 @@ class Edge(SpatialTableModel[EdgeSchema]):
     """
 
     def get_where_edge_type(self, edge_type: str) -> NDArray[np.bool_]:
+        assert self.df is not None
         return (self.df.edge_type == edge_type).to_numpy()
 
     def plot(self, **kwargs) -> Axes:
         ax = kwargs.get("ax", None)
         color_flow = kwargs.get("color_flow", None)
         color_control = kwargs.get("color_control", None)
+        assert self.df is not None
 
         if ax is None:
             _, ax = plt.subplots()
diff --git a/python/ribasim/ribasim/geometry/node.py b/python/ribasim/ribasim/geometry/node.py
index f10a2bc22..5c3752b00 100644
--- a/python/ribasim/ribasim/geometry/node.py
+++ b/python/ribasim/ribasim/geometry/node.py
@@ -80,6 +80,7 @@ def geometry_from_connectivity(
         edge_geometry : np.ndarray
             Array of shapely LineStrings.
         """
+        assert self.df is not None
         geometry = self.df["geometry"]
         from_points = shapely.get_coordinates(geometry.loc[from_id])
         to_points = shapely.get_coordinates(geometry.loc[to_id])
@@ -109,6 +110,7 @@ def connectivity_from_geometry(
         from_node_id : np.ndarray of int
         to_node_id : np.ndarray of int
         """
+        assert self.df is not None
         node_index = self.df.index
         node_xy = shapely.get_coordinates(self.df.geometry.values)
         edge_xy = shapely.get_coordinates(lines)
@@ -143,6 +145,7 @@ def plot_allocation_networks(self, ax=None, zorder=None) -> Any:
 
         contains_main_network = False
         contains_subnetworks = False
+        assert self.df is not None
 
         for allocation_subnetwork_id, df_subnetwork in self.df.groupby(
             "allocation_network_id"
@@ -223,6 +226,7 @@ def plot(self, ax=None, zorder=None) -> Any:
             "User": "g",
             "": "k",
         }
+        assert self.df is not None
 
         for nodetype, df in self.df.groupby("type"):
             assert isinstance(nodetype, str)
@@ -237,6 +241,7 @@ def plot(self, ax=None, zorder=None) -> Any:
                 label=nodetype,
             )
 
+        assert self.df is not None
         geometry = self.df["geometry"]
         for text, xy in zip(self.df.index, np.column_stack((geometry.x, geometry.y))):
             ax.annotate(text=text, xy=xy, xytext=(2.0, 2.0), textcoords="offset points")
diff --git a/python/ribasim/ribasim/model.py b/python/ribasim/ribasim/model.py
index 13b4509df..05057a58c 100644
--- a/python/ribasim/ribasim/model.py
+++ b/python/ribasim/ribasim/model.py
@@ -468,6 +468,8 @@ def print_discrete_control_record(self, path: FilePath) -> None:
         node_attrs, node_instances = zip(*self.nodes().items())
         node_clss = [node_cls.get_input_type() for node_cls in node_instances]
         truth_dict = {"T": ">", "F": "<"}
+        assert self.network.node.df is not None
+        assert self.network.edge.df is not None
 
         if self.discrete_control.condition.df is None:
             raise ValueError("This model has no control input.")