Add config obj: simplification suggestions (#711)

* Removed mode redundancy to simplify the configuration
* Simplify config schema extraction

docs/hooks/generate_readable_schema.py

@@ -14,15 +14,13 @@
 import jsonschema2md
 from mkdocs.structure.files import File
 
-from calliope import AttrDict, config
+from calliope import config
 from calliope.util import schema
 
 TEMPDIR = tempfile.TemporaryDirectory()
 
 SCHEMAS = {
-    "config_schema": AttrDict.from_yaml_string(
-        config.CalliopeConfig().model_yaml_schema()
-    ),
+    "config_schema": config.CalliopeConfig().model_no_ref_schema(),
     "model_schema": schema.MODEL_SCHEMA,
     "math_schema": schema.MATH_SCHEMA,
     "data_table_schema": schema.DATA_TABLE_SCHEMA,

requirements/base.txt

@@ -15,3 +15,4 @@ pyparsing >= 3.0, < 3.1
 ruamel.yaml >= 0.18, < 0.19
 typing-extensions >= 4, < 5
 xarray >= 2024.1, < 2024.4
+pydantic >= 2.9.2

src/calliope/attrdict.py

@@ -329,12 +329,8 @@ def as_dict_flat(self):
             d[k] = self.get_key(k)
         return d
 
-    def to_yaml(self, path=None):
-        """Conversion to YAML.
-
-        Saves the AttrDict to the ``path`` as a YAML file or returns a YAML string
-        if ``path`` is None.
-        """
+    def to_yaml(self, path: str | None = None) -> str:
+        """Return a serialised YAML string."""
         result = self.copy()
         yaml_ = ruamel_yaml.YAML()
         yaml_.indent = 2
@@ -359,13 +355,13 @@ def to_yaml(self, path=None):
         # handle multi-line strings.
         walk_tree(result)
 
-        if path is not None:
+        stream = io.StringIO()
+        yaml_.dump(result, stream)
+        yaml_str = stream.getvalue()
+        if path:
             with open(path, "w") as f:
-                yaml_.dump(result, f)
-        else:
-            stream = io.StringIO()
-            yaml_.dump(result, stream)
-            return stream.getvalue()
+                f.write(yaml_str)
+        return yaml_str
 
     def keys_nested(self, subkeys_as="list"):
         """Returns all keys in the AttrDict, including nested keys.

src/calliope/backend/where_parser.py

@@ -13,16 +13,15 @@
 import xarray as xr
 from typing_extensions import NotRequired, TypedDict
 
+from calliope import config
 from calliope.backend import expression_parser
 from calliope.exceptions import BackendError
 
 if TYPE_CHECKING:
-    from calliope import config
     from calliope.backend.backend_model import BackendModel
 
 
 pp.ParserElement.enablePackrat()
-
 BOOLEANTYPE = np.bool_ | np.typing.NDArray[np.bool_]
 
 

src/calliope/config.py

@@ -5,7 +5,7 @@
 from collections.abc import Hashable
 from datetime import datetime
 from pathlib import Path
-from typing import Annotated, Literal, Self, TypeVar, get_args, overload
+from typing import Annotated, Literal, Self, TypeVar
 
 import jsonref
 from pydantic import AfterValidator, BaseModel, Field, model_validator
@@ -82,27 +82,16 @@ def update(self, update_dict: dict, deep: bool = False) -> Self:
         self._kwargs = update_dict
         return updated
 
-    @overload
-    def model_yaml_schema(self, filepath: str | Path) -> None: ...
-
-    @overload
-    def model_yaml_schema(self, filepath: None = None) -> str: ...
-
-    def model_yaml_schema(self, filepath: str | Path | None = None) -> None | str:
-        """Generate a YAML schema for the class.
-
-        Args:
-            filepath (str | Path | None, optional): If given, save schema to given path. Defaults to None.
+    def model_no_ref_schema(self) -> AttrDict:
+        """Generate an AttrDict with the schema replacing $ref/$def for better readability.
 
         Returns:
-            None | str: If `filepath` is given, returns None. Otherwise, returns the YAML string.
+            AttrDict: class schema.
         """
-        # By default, the schema uses $ref/$def cross-referencing for each pydantic model class,
-        # but this isn't very readable when rendered in our documentation.
-        # So, we resolve references and then delete all the `$defs`
-        schema_dict = AttrDict(jsonref.replace_refs(self.model_json_schema()))
+        schema_dict = AttrDict(super().model_json_schema())
+        schema_dict = AttrDict(jsonref.replace_refs(schema_dict))
         schema_dict.del_key("$defs")
-        return schema_dict.to_yaml(filepath)
+        return schema_dict
 
     @property
     def applied_keyword_overrides(self) -> dict:
@@ -114,21 +103,6 @@ def applied_keyword_overrides(self) -> dict:
         return self._kwargs
 
 
-class ModeBaseModel(ConfigBaseModel):
-    """Mode-specific configuration, which will be hidden from the string representation of the model if that mode is not activated."""
-
-    mode: MODES_T = Field(default="plan")
-    """Mode in which to run the optimisation."""
-
-    @model_validator(mode="after")
-    def update_repr(self) -> Self:
-        """Hide config from model string representation if mode is not activated."""
-        for key, val in self.model_fields.items():
-            if key in get_args(MODES_T):
-                val.repr = self.mode == key
-        return self
-
-
 class Init(ConfigBaseModel):
     """All configuration options used when initialising a Calliope model."""
 
@@ -221,14 +195,18 @@ class BuildOperate(ConfigBaseModel):
     """Which time window to build. This is used to track the window when re-building the model part way through solving in `operate` mode."""
 
 
-class Build(ModeBaseModel):
+class Build(ConfigBaseModel):
     """Base configuration options used when building a Calliope optimisation problem (`calliope.Model.build`)."""
 
     model_config = {
         "title": "build",
         "extra": "allow",
         "revalidate_instances": "always",
     }
+
+    mode: MODES_T = Field(default="plan")
+    """Mode in which to run the optimisation."""
+
     add_math: UniqueList[str] = Field(default=[])
     """
     List of references to files which contain additional mathematical formulations to be applied on top of or instead of the base mode math.
@@ -301,14 +279,13 @@ def require_save_per_spore_path(self) -> Self:
         return self
 
 
-class Solve(ModeBaseModel):
+class Solve(ConfigBaseModel):
     """Base configuration options used when solving a Calliope optimisation problem (`calliope.Model.solve`)."""
 
     model_config = {
         "title": "solve",
         "extra": "forbid",
         "revalidate_instances": "always",
-        "json_schema_extra": hide_from_schema(["mode"]),
     }
 
     save_logs: Path | None = Field(default=None)
@@ -342,28 +319,3 @@ class CalliopeConfig(ConfigBaseModel):
     init: Init = Init()
     build: Build = Build()
     solve: Solve = Solve()
-
-    @model_validator(mode="before")
-    @classmethod
-    def update_solve_mode(cls, data):
-        """Solve mode should match build mode."""
-        data["solve"]["mode"] = data["build"]["mode"]
-        return data
-
-    def update(self, update_dict: dict, deep: bool = False) -> Self:
-        """Return a new iteration of the model with updated fields.
-
-        Updates are validated and stored in the parent class in the `_kwargs` key.
-
-        Args:
-            update_dict (dict): Dictionary with which to update the base model.
-            deep (bool, optional): Set to True to make a deep copy of the model. Defaults to False.
-
-        Returns:
-            BaseModel: New model instance.
-        """
-        update_dict_temp = AttrDict(update_dict)
-        if update_dict_temp.get_key("build.mode", None) is not None:
-            update_dict_temp.set_key("solve.mode", update_dict_temp["build"]["mode"])
-        updated = super().update(update_dict_temp.as_dict(), deep=deep)
-        return updated

src/calliope/model.py

@@ -267,26 +267,26 @@ def build(
             comment="Model: backend build starting",
         )
 
-        this_build_config = self.config.update({"build": kwargs}).build
-        mode = this_build_config.mode
+        build_config = self.config.update({"build": kwargs}).build
+        mode = build_config.mode
         if mode == "operate":
             if not self._model_data.attrs["allow_operate_mode"]:
                 raise exceptions.ModelError(
                     "Unable to run this model in operate (i.e. dispatch) mode, probably because "
                     "there exist non-uniform timesteps (e.g. from time clustering)"
                 )
-            backend_input = self._prepare_operate_mode_inputs(this_build_config.operate)
+            backend_input = self._prepare_operate_mode_inputs(build_config.operate)
         else:
             backend_input = self._model_data
 
-        init_math_list = [] if this_build_config.ignore_mode_math else [mode]
+        init_math_list = [] if build_config.ignore_mode_math else [mode]
         end_math_list = [] if add_math_dict is None else [add_math_dict]
-        full_math_list = init_math_list + this_build_config.add_math + end_math_list
+        full_math_list = init_math_list + build_config.add_math + end_math_list
         LOGGER.debug(f"Math preprocessing | Loading math: {full_math_list}")
         model_math = preprocess.CalliopeMath(full_math_list, self.config.init.def_path)
 
         self.backend = backend.get_model_backend(
-            this_build_config, backend_input, model_math
+            build_config, backend_input, model_math
         )
         self.backend.add_optimisation_components()
 
@@ -341,26 +341,24 @@ def solve(self, force: bool = False, warmstart: bool = False, **kwargs) -> None:
         else:
             to_drop = []
 
-        kwargs["mode"] = self.config.build.applied_keyword_overrides.get(
-            "mode", self.config.build.mode
-        )
-
-        this_solve_config = self.config.update({"solve": kwargs}).solve
+        solve_config = self.config.update({"solve": kwargs}).solve
+        # FIXME: find a way to avoid overcomplicated passing of settings between modes
+        mode = self.config.update(self.config.applied_keyword_overrides).build.mode
         self._model_data.attrs["timestamp_solve_start"] = log_time(
             LOGGER,
             self._timings,
             "solve_start",
-            comment=f"Optimisation model | starting model in {this_solve_config.mode} mode.",
+            comment=f"Optimisation model | starting model in {mode} mode.",
         )
 
-        shadow_prices = this_solve_config.shadow_prices
+        shadow_prices = solve_config.shadow_prices
         self.backend.shadow_prices.track_constraints(shadow_prices)
 
-        if this_solve_config.mode == "operate":
-            results = self._solve_operate(**this_solve_config.model_dump())
+        if mode == "operate":
+            results = self._solve_operate(**solve_config.model_dump())
         else:
             results = self.backend._solve(
-                warmstart=warmstart, **this_solve_config.model_dump()
+                warmstart=warmstart, **solve_config.model_dump()
             )
 
         log_time(