Merge pull request #104 from ImperialCollegeLondon/extensions_patch

✨ Add load_extensions_files step

tests/test_model.py

@@ -4,10 +4,10 @@
 @author: Barney
 
 """
-
-# import pytest
+import os
+import pytest
 import unittest
-from unittest import TestCase
+from unittest import TestCase, mock
 
 from wsimod.arcs.arcs import Arc
 from wsimod.nodes.land import Land
@@ -305,5 +305,47 @@ def test_customise_orchestration(self):
         self.assertListEqual(my_model.orchestration, revised_orchestration)
 
 
+class TestLoadExtensionFiles:
+    def test_load_extension_files_valid(self, tmp_path_factory):
+        from wsimod.orchestration.model import load_extension_files
+
+        with tmp_path_factory.mktemp("extensions") as tempdir:
+            valid_file = os.path.join(tempdir, "valid_extension.py")
+            with open(valid_file, "w") as f:
+                f.write("def test_func(): pass")
+
+            load_extension_files([valid_file])
+
+    def test_load_extension_files_invalid_extension(self, tmp_path_factory):
+        from wsimod.orchestration.model import load_extension_files
+
+        with tmp_path_factory.mktemp("extensions") as tempdir:
+            invalid_file = os.path.join(tempdir, "invalid_extension.txt")
+            with open(invalid_file, "w") as f:
+                f.write("This is a text file")
+
+            with pytest.raises(ValueError, match="Only .py files are supported"):
+                load_extension_files([invalid_file])
+
+    def test_load_extension_files_nonexistent_file(self):
+        from wsimod.orchestration.model import load_extension_files
+
+        with pytest.raises(
+            FileNotFoundError, match="File nonexistent_file.py does not exist"
+        ):
+            load_extension_files(["nonexistent_file.py"])
+
+    def test_load_extension_files_import_error(self, tmp_path_factory):
+        from wsimod.orchestration.model import load_extension_files
+
+        with tmp_path_factory.mktemp("extensions") as tempdir:
+            valid_file = os.path.join(tempdir, "valid_extension.py")
+            with open(valid_file, "w") as f:
+                f.write("raise ImportError")
+
+            with pytest.raises(ImportError):
+                load_extension_files([valid_file])
+
+
 if __name__ == "__main__":
     unittest.main()

wsimod/orchestration/model.py

@@ -178,7 +178,7 @@ def load(self, address, config_name="config.yml", overrides={}):
         from ..extensions import apply_patches
 
         with open(os.path.join(address, config_name), "r") as file:
-            data = yaml.safe_load(file)
+            data: dict = yaml.safe_load(file)
 
         for key, item in overrides.items():
             data[key] = item
@@ -220,6 +220,7 @@ def load(self, address, config_name="config.yml", overrides={}):
         if "dates" in data.keys():
             self.dates = [to_datetime(x) for x in data["dates"]]
 
+        load_extension_files(data.get("extensions", []))
         apply_patches(self)
 
     def save(self, address, config_name="config.yml", compress=False):
@@ -1269,3 +1270,27 @@ def yaml2csv(address, config_name="config.yml", csv_folder_name="csv"):
                     writer.writerow(
                         [str(value_[x]) if x in value_.keys() else None for x in fields]
                     )
+
+
+def load_extension_files(files: list[str]) -> None:
+    """Load extension files from a list of files.
+
+    Args:
+        files (list[str]): List of file paths to load
+
+    Raises:
+        ValueError: If file is not a .py file
+        FileNotFoundError: If file does not exist
+    """
+    import importlib
+    from pathlib import Path
+
+    for file in files:
+        if not file.endswith(".py"):
+            raise ValueError(f"Only .py files are supported. Invalid file: {file}")
+        if not Path(file).exists():
+            raise FileNotFoundError(f"File {file} does not exist")
+
+        spec = importlib.util.spec_from_file_location("module.name", file)
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)