From d4bcad6684663c5f5fceb791f7949c2cd445e95d Mon Sep 17 00:00:00 2001
From: Hofer-Julian <julianhofer@gnome.org>
Date: Tue, 13 Feb 2024 10:04:34 +0100
Subject: [PATCH] Improve Node representation

---
 python/ribasim/playground.py  | 18 +-----------------
 python/ribasim/ribasim/add.py | 22 ++++++++--------------
 2 files changed, 9 insertions(+), 31 deletions(-)

diff --git a/python/ribasim/playground.py b/python/ribasim/playground.py
index fd6fa3c78..68011a634 100644
--- a/python/ribasim/playground.py
+++ b/python/ribasim/playground.py
@@ -1,21 +1,5 @@
 # %%
-from pydantic import BaseModel
-
-
-class A(BaseModel):
-    x: float
-    y: float
-
-
-class B(BaseModel):
-    a: A
-
-
-b = B(a={"x": 1.2, "y": 1.4})
-
-print(b.model_dump())
-
-# %%
+import geopandas
 import pandas as pd
 from pandas import DataFrame
 from ribasim.add import Basins
diff --git a/python/ribasim/ribasim/add.py b/python/ribasim/ribasim/add.py
index 404e216cf..00cea573f 100644
--- a/python/ribasim/ribasim/add.py
+++ b/python/ribasim/ribasim/add.py
@@ -5,6 +5,7 @@
 import pandas as pd
 from pandas import DataFrame, Timestamp
 from pydantic import ConfigDict, Field
+from shapely.geometry import Point
 
 from ribasim.config import Basin
 from ribasim.geometry.node import NodeSchema
@@ -37,7 +38,7 @@
 
 
 class _BaseModel(BaseModel):
-    model_config = ConfigDict(extra="forbid")
+    model_config = ConfigDict(extra="forbid", validate_assignment=True)
 
 
 class Node(_BaseModel):
@@ -68,9 +69,7 @@ class Basins(_BaseModel):
         default_factory=TableModel[BasinSubgridSchema],
         json_schema_extra={"sort_keys": ["subgrid_id", "basin_level"]},
     )
-    node: SpatialTableModel[NodeSchema] = Field(
-        default_factory=SpatialTableModel[NodeSchema],
-    )
+    node: list[Node] = []
 
     def __init__(self, basins: list[Basin]):
         super().__init__()
@@ -78,17 +77,12 @@ def __init__(self, basins: list[Basin]):
             id = basin.node["id"]
 
             basin_dict = basin.model_dump()
+            node = basin_dict.pop("node")
+            self.node.append(Node(**node))
             for basin_key, basin_value in basin_dict.items():
-                if basin_key == "node":
-                    node_dict = {
-                        node_key: [value] for node_key, value in basin_value.items()
-                    }
-                    table_to_append = pd.DataFrame(node_dict)
-                else:
-                    table_to_append = basin_value.assign(node_id=id)
-
-                original_table = getattr(self, basin_key)
-                original_table.df = pd.concat([original_table.df, table_to_append])
+                member = getattr(self, basin_key)
+                table_to_append = basin_value.assign(node_id=id)
+                member.df = pd.concat([member.df, table_to_append])
 
 
 class Model(_BaseModel):