Updated model API

src/ewatercycle/base/model.py

@@ -2,9 +2,10 @@
 
 import abc
 import logging
-from datetime import datetime, timezone
+import datetime
+from datetime import timezone
 from pathlib import Path
-from typing import Any, Iterable, Type, cast
+from typing import Any, Iterable, Optional, Type, cast
 
 import bmipy
 import numpy as np
@@ -86,19 +87,20 @@ def setup(self, *, cfg_dir: str | None = None, **kwargs) -> tuple[str, str]:
 
         return str(self._cfg_file), str(self._cfg_dir)
 
-    def _make_cfg_dir(self, cfg_dir):
+    def _make_cfg_dir(self, cfg_dir: Optional[str] = None, **kwargs) -> Path:
         if cfg_dir is not None:
-            cfg_dir = to_absolute_path(cfg_dir)
+            cfg_path = to_absolute_path(cfg_dir)
         else:
             tz = timezone.utc
-            timestamp = datetime.now(tz).strftime("%Y%m%d_%H%M%S")
-            cfg_dir = to_absolute_path(
-                f"ewatercycle_{timestamp}", parent=CFG.output_dir
+            timestamp = datetime.datetime.now(tz).strftime("%Y%m%d_%H%M%S")
+            folder_prefix = kwargs.get("folder_prefix", "ewatercycle")
+            cfg_path = to_absolute_path(
+                f"{folder_prefix}_{timestamp}", parent=CFG.output_dir
             )
 
-        cfg_dir.mkdir(parents=True, exist_ok=True)
+        cfg_path.mkdir(parents=True, exist_ok=True)
 
-        return cfg_dir
+        return cfg_path
 
     def _make_cfg_file(self, **kwargs):
         """Create new config file and return its path."""
@@ -312,7 +314,7 @@ def time_as_isostr(self) -> str:
         return self.time_as_datetime.strftime(ISO_TIMEFMT)
 
     @property
-    def start_time_as_datetime(self) -> datetime:
+    def start_time_as_datetime(self) -> datetime.datetime:
         """Start time of the model as a datetime object."""
         return num2date(
             self._bmi.get_start_time(),
@@ -321,7 +323,7 @@ def start_time_as_datetime(self) -> datetime:
         )
 
     @property
-    def end_time_as_datetime(self) -> datetime:
+    def end_time_as_datetime(self) -> datetime.datetime:
         """End time of the model as a datetime object'."""
         return num2date(
             self._bmi.get_end_time(),
@@ -330,7 +332,7 @@ def end_time_as_datetime(self) -> datetime:
         )
 
     @property
-    def time_as_datetime(self) -> datetime:
+    def time_as_datetime(self) -> datetime.datetime:
         """Current time of the model as a datetime object'."""
         return num2date(
             self._bmi.get_current_time(),

src/ewatercycle/plugins/hype/model.py

@@ -12,12 +12,11 @@
 from dateutil.tz import UTC
 from pydantic import PrivateAttr, root_validator
 
-from ewatercycle import CFG
 from ewatercycle.base.model import ISO_TIMEFMT, ContainerizedModel
 from ewatercycle.base.parameter_set import ParameterSet
 from ewatercycle.container import ContainerImage
 from ewatercycle.plugins.hype.forcing import HypeForcing
-from ewatercycle.util import geographical_distances, get_time, to_absolute_path
+from ewatercycle.util import geographical_distances, get_time
 
 logger = logging.getLogger(__name__)
 
@@ -137,29 +136,33 @@ def _export_config(self) -> Path:
 
         return cfg_file
 
-    def _make_cfg_dir(self, cfg_dir: Optional[Path] = None) -> Path:
-        if cfg_dir:
-            cfg_dir = to_absolute_path(cfg_dir)
-        else:
-            # Must exist before setting up default config
-            timestamp = datetime.datetime.now(datetime.timezone.utc).strftime(
-                "%Y%m%d_%H%M%S"
-            )
-            cfg_dir = to_absolute_path(f"hype_{timestamp}", parent=CFG.output_dir)
-        cfg_dir.mkdir(parents=True, exist_ok=True)
+    def _make_cfg_dir(
+        self,
+        cfg_dir: Optional[str] = None,
+        **kwargs,
+    ) -> Path:
+        """Make sure there is a working directory.
+
+        Args:
+            cfg_dir: If cfg dir is None or does not exist then create sub-directory
+                in CFG.output_dir
+        """
+        cfg_path = super()._make_cfg_dir(
+            cfg_dir=cfg_dir, folder_prefix="hype", **kwargs
+        )
 
         # copy parameter set files to cfg_dir
         assert self.parameter_set
         shutil.copytree(
-            src=self.parameter_set.directory, dst=cfg_dir, dirs_exist_ok=True
+            src=self.parameter_set.directory, dst=cfg_path, dirs_exist_ok=True
         )
 
         # copy forcing files to cfg_dir
         if self.forcing is not None and self.forcing.directory is not None:
             forcing_dir = self.forcing.directory
-            shutil.copytree(src=forcing_dir, dst=cfg_dir, dirs_exist_ok=True)
+            shutil.copytree(src=forcing_dir, dst=cfg_path, dirs_exist_ok=True)
 
-        return cfg_dir
+        return cfg_path
 
     def _make_bmi_instance(self) -> bmipy.Bmi:
         """Make the bmi instance and overwrite 'get_time_units' method."""

src/ewatercycle/plugins/pcrglobwb/model.py

@@ -1,13 +1,12 @@
 """eWaterCycle wrapper around PCRGlobWB BMI."""
 
-import datetime
 import logging
 from os import PathLike
+from pathlib import Path
 from typing import Optional
 
 from pydantic import PrivateAttr, root_validator
 
-from ewatercycle import CFG
 from ewatercycle.base.model import ContainerizedModel
 from ewatercycle.base.parameter_set import ParameterSet
 from ewatercycle.container import ContainerImage
@@ -91,18 +90,20 @@ def _update_parameters(cls, values):
         )
         return values
 
-    def _setup_work_dir(self, cfg_dir: Optional[str] = None):
-        if cfg_dir:
-            self.work_dir = to_absolute_path(cfg_dir)
-        else:
-            # Must exist before setting up default config
-            timestamp = datetime.datetime.now(datetime.timezone.utc).strftime(
-                "%Y%m%d_%H%M%S"
-            )
-            self.work_dir = to_absolute_path(
-                f"pcrglobwb_{timestamp}", parent=CFG.output_dir
-            )
-        self.work_dir.mkdir(parents=True, exist_ok=True)
+    def _make_cfg_dir(
+        self,
+        cfg_dir: Optional[str] = None,
+        **kwargs,
+    ) -> Path:
+        """Make sure there is a working directory.
+
+        Args:
+            cfg_dir: If cfg dir is None or does not exist then create sub-directory
+                in CFG.output_dir
+        """
+        return super()._make_cfg_dir(
+            cfg_dir=cfg_dir, folder_prefix="pcrglobwb", **kwargs
+        )
 
     def _make_cfg_file(self, **kwargs):
         self._update_config(**kwargs)
@@ -143,9 +144,9 @@ def _update_config(self, **kwargs):
             )
 
     def _export_config(self) -> PathLike:
-        self._config.set("globalOptions", "outputDir", str(self.work_dir))
+        self._config.set("globalOptions", "outputDir", str(self._cfg_dir))
         new_cfg_file = to_absolute_path(
-            "pcrglobwb_ewatercycle.ini", parent=self.work_dir
+            "pcrglobwb_ewatercycle.ini", parent=self._cfg_dir
         )
         with new_cfg_file.open("w") as filename:
             self._config.write(filename)

src/ewatercycle/plugins/wflow/model.py

@@ -94,7 +94,7 @@ def _make_cfg_file(self, **kwargs) -> Path:
 
         return cfg_file
 
-    def _make_cfg_dir(self, cfg_dir: Optional[str | Path] = None) -> Path:
+    def _make_cfg_dir(self, cfg_dir: Optional[str] = None, **kwargs) -> Path:
         """Create working directory for parameter sets, forcing and wflow config."""
         if cfg_dir:
             self._work_dir = to_absolute_path(cfg_dir)

tests/plugins/hype/test_model.py

@@ -83,9 +83,12 @@ def model(self, parameter_set):
 
     @pytest.fixture
     def model_with_setup(self, mocked_config, model: Hype):
-        with patch.object(
+        with (
+            patch.object(
             BmiClientApptainer, "__init__", return_value=None
-        ) as mocked_constructor, patch("datetime.datetime") as mocked_datetime:
+            ) as mocked_constructor, 
+            patch("datetime.datetime") as mocked_datetime
+        ):
             mocked_datetime.now.return_value = datetime(2021, 1, 2, 3, 4, 5)
             config_file, config_dir = model.setup()
         return config_file, config_dir, mocked_constructor