Apply overrides for data_input_dict in nodes

ImperialCollegeLondon · Mar 18, 2024 · 3b93e99 · 3b93e99
1 parent 5a26781
commit 3b93e99
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 1 deletion.
diff --git a/tests/example_data_input_dict.csv.gz b/tests/example_data_input_dict.csv.gz
diff --git a/tests/test_nodes.py b/tests/test_nodes.py
@@ -21,6 +21,7 @@
 )
 from wsimod.nodes.storage import Storage
 from wsimod.nodes.waste import Waste
+import os
 
 
 class MyTestClass(TestCase):
@@ -659,7 +660,22 @@ def test_decay_queue(self):
             "temperature": ((5 * 10) + (15 * 1)) / 6,
         }
         self.assertDictAlmostEqual(d4, tank.storage, 15)
-
+
+    def test_node_overrides(self):
+        node = Node(name="", data_input_dict={("temperature", 1): 15})
+        # check the format of dict
+        new_data_input_dict = {("temperature", 1): 10,
+                               ("temperature", 2): 20,
+                               }
+        node.apply_overrides({'data_input_dict': new_data_input_dict})
+        self.assertDictEqual(node.data_input_dict, new_data_input_dict)
+        # check the format of str
+        new_data_input_dict = "example_data_input_dict.csv.gz"
+        node.apply_overrides({'data_input_dict': new_data_input_dict})
+        from wsimod.orchestration.model import read_csv
+        new_data_input_dict = read_csv(new_data_input_dict)
+        self.assertDictEqual(node.data_input_dict, new_data_input_dict)
+        print(dict(list(node.data_input_dict.items())[:5]))
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/wsimod/nodes/nodes.py b/wsimod/nodes/nodes.py
@@ -88,6 +88,16 @@ def apply_overrides(self, overrides: Dict[str, Any] = {}) -> None:
         Args:
             overrides (dict, optional): Dictionary of overrides. Defaults to {}.
         """
+        # overrides data_input_dict
+        from wsimod.orchestration.model import read_csv
+        content = overrides.pop('data_input_dict', self.data_input_dict)
+        if isinstance(content, str):
+            self.data_input_dict = read_csv(content)
+        elif isinstance(content, dict):
+            self.data_input_dict = content
+        else:
+            print('ERROR: not recognised format for data_input_dict')
+
         if len(overrides) > 0:
             print(f"No override behaviour defined for: {overrides.keys()}")