Move to_yaml from AttrDict to io

Author: irm-codebase

src/calliope/attrdict.py

@@ -3,11 +3,8 @@
 """AttrDict implementation (a subclass of regular dict) used for managing model configuration."""
 
 import copy
-import io
 import logging
 
-import numpy as np
-import ruamel.yaml as ruamel_yaml
 from typing_extensions import Self
 
 logger = logging.getLogger(__name__)
@@ -187,41 +184,6 @@ 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.
-        """
-        result = self.copy()
-        yaml_ = ruamel_yaml.YAML()
-        yaml_.indent = 2
-        yaml_.block_seq_indent = 0
-        yaml_.sort_base_mapping_type_on_output = False
-
-        # Numpy objects should be converted to regular Python objects,
-        # so that they are properly displayed in the resulting YAML output
-        for k in result.keys_nested():
-            # Convert numpy numbers to regular python ones
-            v = result.get_key(k)
-            if isinstance(v, np.floating):
-                result.set_key(k, float(v))
-            elif isinstance(v, np.integer):
-                result.set_key(k, int(v))
-            # Lists are turned into seqs so that they are formatted nicely
-            elif isinstance(v, list):
-                result.set_key(k, yaml_.seq(v))
-
-        result = result.as_dict()
-
-        if path is not None:
-            with open(path, "w") as f:
-                yaml_.dump(result, f)
-        else:
-            stream = io.StringIO()
-            yaml_.dump(result, stream)
-            return stream.getvalue()
-
     def keys_nested(self, subkeys_as="list"):
         """Returns all keys in the AttrDict, including nested keys.
 

src/calliope/backend/backend_model.py

@@ -30,7 +30,7 @@
 from calliope.attrdict import AttrDict
 from calliope.backend import helper_functions, parsing
 from calliope.exceptions import warn as model_warn
-from calliope.io import load_config
+from calliope.io import load_config, to_yaml
 from calliope.preprocess.model_math import ORDERED_COMPONENTS_T, CalliopeMath
 from calliope.util.schema import (
     MODEL_SCHEMA,
@@ -449,7 +449,7 @@ def _add_to_dataset(
                 yaml_snippet_attrs[attr] = val
 
         if yaml_snippet_attrs:
-            add_attrs["yaml_snippet"] = AttrDict(yaml_snippet_attrs).to_yaml()
+            add_attrs["yaml_snippet"] = to_yaml(yaml_snippet_attrs)
 
         da.attrs = {
             "obj_type": obj_type,

src/calliope/cli.py

@@ -13,7 +13,7 @@
 
 import click
 
-from calliope import AttrDict, Model, examples, read_netcdf
+from calliope import Model, examples, io, read_netcdf
 from calliope._version import __version__
 from calliope.exceptions import BackendError
 from calliope.util.generate_runs import generate
@@ -400,4 +400,4 @@ def generate_scenarios(
             }
         }
 
-        AttrDict(scenarios).to_yaml(out_file)
+        io.to_yaml(scenarios, path=out_file)

src/calliope/io.py

@@ -6,6 +6,7 @@
 import logging
 import os
 from copy import deepcopy
+from io import StringIO
 from pathlib import Path
 
 # We import netCDF4 before xarray to mitigate a numpy warning:
@@ -23,6 +24,8 @@
 logger = logging.getLogger(__name__)
 
 CONFIG_DIR = importlib.resources.files("calliope") / "config"
+YAML_INDENT = 2
+YAML_BLOCK_SEQUENCE_INDENT = 0
 
 
 def read_netcdf(path):
@@ -75,7 +78,7 @@ def _serialise(attrs: dict) -> None:
     dict_attrs = [k for k, v in attrs.items() if isinstance(v, dict)]
     attrs["serialised_dicts"] = dict_attrs
     for attr in dict_attrs:
-        attrs[attr] = AttrDict(attrs[attr]).to_yaml()
+        attrs[attr] = to_yaml(attrs[attr])
 
     # Convert boolean attrs to ints
     bool_attrs = [k for k, v in attrs.items() if isinstance(v, bool)]
@@ -287,3 +290,42 @@ def _resolve_yaml_imports(
         loaded_dict.del_key("import")
 
     return loaded_dict
+
+
+def to_yaml(data: AttrDict | dict, path: None | str | Path = None) -> str:
+    """Conversion to YAML.
+
+    Saves the AttrDict to the ``path`` as a YAML file or returns a YAML string
+    if ``path`` is None.
+    """
+    result = AttrDict(data).copy()
+    # Prepare YAML parsing settings
+    yaml_ = ruamel_yaml.YAML()
+    yaml_.indent = YAML_INDENT
+    yaml_.block_seq_indent = YAML_BLOCK_SEQUENCE_INDENT
+    yaml_.sort_base_mapping_type_on_output = (
+        False  # FIXME: identify if this is necessary
+    )
+
+    # Numpy objects should be converted to regular Python objects,
+    # so that they are properly displayed in the resulting YAML output
+    for k in result.keys_nested():
+        # Convert numpy numbers to regular python ones
+        v = result.get_key(k)
+        if isinstance(v, np.floating):
+            result.set_key(k, float(v))
+        elif isinstance(v, np.integer):
+            result.set_key(k, int(v))
+        # Lists are turned into seqs so that they are formatted nicely
+        elif isinstance(v, list):
+            result.set_key(k, yaml_.seq(v))
+
+    result = result.as_dict()
+
+    if path is not None:
+        with open(path, "w") as f:
+            yaml_.dump(result, f)
+
+    stream = StringIO()
+    yaml_.dump(result, stream)
+    return stream.getvalue()

src/calliope/preprocess/model_definition.py

@@ -7,7 +7,7 @@
 
 from calliope import exceptions
 from calliope.attrdict import AttrDict
-from calliope.io import read_rich_yaml
+from calliope.io import read_rich_yaml, to_yaml
 from calliope.util.tools import listify
 
 LOGGER = logging.getLogger(__name__)
@@ -129,7 +129,7 @@ def _combine_overrides(overrides: AttrDict, scenario_overrides: list):
     combined_override_dict = AttrDict()
     for override in scenario_overrides:
         try:
-            yaml_string = overrides[override].to_yaml()
+            yaml_string = to_yaml(overrides[override])
             override_with_imports = read_rich_yaml(yaml_string)
         except KeyError:
             raise exceptions.ModelError(f"Override `{override}` is not defined.")

tests/test_core_attrdict.py

@@ -206,7 +206,3 @@ def test_del_key_single(self, attr_dict):
     def test_del_key_nested(self, attr_dict):
         attr_dict.del_key("c.z.I")
         assert "I" not in attr_dict.c.z
-
-    def test_to_yaml_string(self, attr_dict):
-        result = attr_dict.to_yaml()
-        assert "a: 1" in result