Added get_source_type

HISTORY.rst

@@ -4,7 +4,7 @@ History
 2.3.2 (unreleased)
 ------------------
 
-- Nothing changed yet.
+- Added get_source_type() to GridH5StructureControl.
 
 
 2.3.1 (2024-10-21)

tests/test_structure_controls.py

@@ -8,7 +8,10 @@
 from threedigrid.admin.structure_controls.exporters import (
     structure_control_actions_to_csv,
 )
-from threedigrid.admin.structure_controls.models import StructureControl
+from threedigrid.admin.structure_controls.models import (
+    StructureControl,
+    StructureControlSourceTypes,
+)
 
 test_file_dir = os.path.join(
     os.path.dirname(os.path.abspath(__file__)), "test_files/structure_controls"
@@ -104,3 +107,20 @@ def test_export_method(gsc, tmp_path):
     csv_path = os.path.join(tmp_path, "test_struct_control_actions.csv")
     structure_control_actions_to_csv(gsc, csv_path)
     assert os.path.exists(csv_path)
+
+
+def test_get_source_type(gsc: GridH5StructureControl):
+    struct_cntrls: List[StructureControl] = gsc.table_control.group_by_action_type(
+        "set_pump_capacity"
+    )
+    assert len(struct_cntrls) == 2
+    assert isinstance(struct_cntrls[0], StructureControl)
+    assert struct_cntrls[0].source_type == StructureControlSourceTypes.PUMPS
+    assert struct_cntrls[1].source_type == StructureControlSourceTypes.PUMPS
+
+
+def test_get_source_type_2(gsc: GridH5StructureControl):
+    assert gsc.get_source_type("set_pump_capacity") == StructureControlSourceTypes.PUMPS
+    assert gsc.get_source_type("set_crest_level") == StructureControlSourceTypes.LINES
+    with pytest.raises(NotImplementedError):
+        gsc.get_source_type("this_action_most_likely_does_not_exist")

threedigrid/admin/gridresultadmin.py

@@ -41,6 +41,7 @@
 )
 from threedigrid.admin.structure_controls.models import (
     StructureControl,
+    StructureControlSourceTypes,
     StructureControlTypes,
 )
 from threedigrid.orm.models import Model
@@ -355,6 +356,19 @@ def get_source_table(self, action_type, grid_id):
 
         return source_table, source_table_id
 
+    def get_source_type(self, action_type):
+        """Indicates whether the action is applied on a line, pump (or node) feature"""
+        if action_type == "set_pump_capacity":
+            return StructureControlSourceTypes.PUMPS
+        elif action_type in [
+            "set_discharge_coefficients",
+            "set_crest_level",
+            "set_gate_level",
+        ]:
+            return StructureControlSourceTypes.LINES
+        else:
+            raise NotImplementedError
+
 
 class _GridH5NestedStructureControl:
     def __init__(
@@ -447,12 +461,14 @@ def group_by_id(self, id: str) -> Optional[StructureControl]:
         source_table, source_table_id = self.struct_control.get_source_table(
             action_type, grid_id
         )
+        source_type = self.struct_control.get_source_type(action_type)
 
         return StructureControl(
             id=id,
             grid_id=grid_id,
             source_table=source_table,
             source_table_id=source_table_id,
+            source_type=source_type,
             time=self.time[mask],
             action_type=action_type,
             action_value_1=self.action_value_1[mask],

threedigrid/admin/structure_controls/models.py

@@ -1,4 +1,4 @@
-from enum import Enum
+from enum import Enum, StrEnum
 from typing import List
 
 
@@ -8,13 +8,20 @@ class StructureControlTypes(Enum):
     timed_control = "timed_control"
 
 
+class StructureControlSourceTypes(StrEnum):
+    LINES = "lines"
+    NODES = "nodes"
+    PUMPS = "pumps"
+
+
 class StructureControl:
     def __init__(
         self,
         id: str,
         grid_id: int,
         source_table: str,
         source_table_id: int,
+        source_type: StructureControlSourceTypes,
         time: List[float],
         action_type: str,
         action_value_1: List[float],
@@ -25,6 +32,7 @@ def __init__(
         self.grid_id = grid_id
         self.source_table = source_table
         self.source_table_id = source_table_id
+        self.source_type = source_type
         self.time = time
         self.action_type = action_type
         self.action_value_1 = action_value_1

threedigrid/orm/base/models.py

@@ -572,5 +572,5 @@ def __repr__(self):
         return "<orm {} instance of {}>".format(self.__contenttype__(), self.model_name)
 
     def __contenttype__(self):
-        """conent type, e.g. lines, nodes, cells, ..."""
+        """content type, e.g. lines, nodes, cells, ..."""
         return self._datasource.group_name