Automatically name index

Deltares · Dec 13, 2024 · 6d2ea71 · 6d2ea71
1 parent 5c08e2f
commit 6d2ea71
Show file tree

Hide file tree

Showing 7 changed files with 191 additions and 14 deletions.
diff --git a/python/ribasim/ribasim/geometry/edge.py b/python/ribasim/ribasim/geometry/edge.py
@@ -49,9 +49,10 @@ class EdgeSchema(_GeoBaseSchema):
     edge_type: Series[str] = pa.Field(default="flow")
     geometry: GeoSeries[LineString] = pa.Field(default=None, nullable=True)
 
-    @classmethod
-    def _index_name(self) -> str:
-        return "edge_id"
+    @pa.dataframe_parser
+    def _name_index(cls, df):
+        df.index.name = "edge_id"
+        return df
 
 
 class EdgeTable(SpatialTableModel[EdgeSchema]):

diff --git a/python/ribasim/ribasim/geometry/node.py b/python/ribasim/ribasim/geometry/node.py
@@ -27,9 +27,10 @@ class NodeSchema(_GeoBaseSchema):
     )
     geometry: GeoSeries[Point] = pa.Field(default=None, nullable=True)
 
-    @classmethod
-    def _index_name(self) -> str:
-        return "node_id"
+    @pa.dataframe_parser
+    def _name_index(cls, df):
+        df.index.name = "node_id"
+        return df
 
 
 class NodeTable(SpatialTableModel[NodeSchema]):

diff --git a/python/ribasim/ribasim/input_base.py b/python/ribasim/ribasim/input_base.py
@@ -217,7 +217,6 @@ def _check_dataframe(cls, value: Any) -> Any:
 
         # Enable initialization with a DataFrame.
         if isinstance(value, pd.DataFrame | gpd.GeoDataFrame):
-            value.index.rename("fid", inplace=True)
             value = {"df": value}
 
         return value
@@ -386,7 +385,6 @@ def _from_db(cls, path: Path, table: str):
                     # tell pyarrow to map to pd.ArrowDtype rather than NumPy
                     arrow_to_pandas_kwargs={"types_mapper": pd.ArrowDtype},
                 )
-                df.index.rename(cls.tableschema()._index_name(), inplace=True)
             else:
                 df = None